__precompile__()
module Decomp

using StaticArrays: MVector, SVector
export eigen!, eigen, inphdecomp, propdecomp, inphmodel, propmodel, mixmodel

 function eigen!(e11::T,e22::T,e33::T,e12::T,e13::T,e23::T,eig::MVector{3,T},eigv1::MVector{3,T},eigv2::MVector{3,T},eigv3::MVector{3,T}) where T<:AbstractFloat

  p1 = muladd(e12, e12, muladd(e13, e13, e23*e23))
  if p1 == 0.0
    # E is diagonal.
    # eig[1] = e11
    # eig[2] = e22
    # eig[3] = e33

    if abs(e11) >= abs(e22)
      if abs(e11) >= abs(e33)
        eig[1]=e11

        eigv1[1] = oneunit(T)
        eigv1[2] = zero(T)
        eigv1[3] = zero(T)

        if abs(e22) >= abs(e33)
          eig[2] = e22
          eig[3] = e33

          eigv2[1] = zero(T)
          eigv2[2] = oneunit(T)
          eigv2[3] = zero(T)
        else
          eig[3] = e22
          eig[2] = e33

          eigv2[1] = zero(T)
          eigv2[2] = zero(T)
          eigv2[3] = oneunit(T)
        end
      else
        eig[1] = e33
        eig[2] = e11
        eig[3] = e22

        eigv1[1] = zero(T)
        eigv1[2] = zero(T)
        eigv1[3] = oneunit(T)

        eigv2[1] = oneunit(T)
        eigv2[2] = zero(T)
        eigv2[3] = zero(T)
      end
    else
      if abs(e22) >= abs(e33)
        eig[1] = e22

        eigv1[1] = zero(T)
        eigv1[2] = oneunit(T)
        eigv1[3] = zero(T)

      elseif abs(e33) >= abs(e11)
        eig[2] = e33
        eig[3] = e11

        eigv2[1] = zero(T)
        eigv2[2] = zero(T)
        eigv2[3] = oneunit(T)

      else
        eig[3] = e33
        eig[2] = e11

        eigv2[1] = oneunit(T)
        eigv2[2] = zero(T)
        eigv2[3] = zero(T)
      end
    end

    eigv3[1] = eigv1[2]*eigv2[3] - (eigv1[3]*eigv2[2])
    eigv3[2] = eigv1[3]*eigv2[1] - (eigv1[1]*eigv2[3])
    eigv3[3] = eigv1[1]*eigv2[2] - (eigv1[2]*eigv2[1])
  else
    q = (e11 + e22 + e33)/3
    p2 = (e11-q)^2 + (e22-q)^2 + (e33-q)^2 + 2*p1
    p = sqrt(p2/6)
    r = ((e11-q)*(e22-q)*(e33-q) - (e11-q)*(e23^2) - (e12^2)*(e33-q) + 2*(e12*e13*e23) - (e13^2)*(e22-q))/(2*p*p*p)

    # In exact arithmetic for a symmetric matrix  -1 <= r <= 1
    # but computation error can leave it slightly outside this range.
    if r <= -1
      ϕ = T(π/3)
    elseif r >= 1
      ϕ = zero(T)
    else
      ϕ = acos(r)/3
    end

    # the eigenvalues satisfy eig[3] <= eig[2] <= eig[1]
    eig[1] = q + 2*p*cos(ϕ)
    eig[3] = q + 2*p*cos(ϕ+(2*π/3))
    eig[2] = 3*q - eig[1] - eig[3]     # since trace(E) = eig[1] + eig[2] + eig[3] = 3q

    #sort!(eig,by=abs,rev=true)
    if abs(eig[1]) < abs(eig[3])
      big = eig[3]
      if abs(eig[1]) < abs(eig[2])
        eig[3] = eig[1]
        eig[1] = big
      else
        eig[3] = eig[2]
        eig[2] = eig[1]
        eig[1] = big
      end
    else
      if abs(eig[2]) < abs(eig[3])
        medium = eig[3]
        eig[3] = eig[2]
        eig[2] = medium
      end
    end

    bla = ((e22 - eig[1])*(e33 - eig[1]) - e23*e23)
    if bla != 0
      eigv1[1] = oneunit(T)
      eigv1[2] = (e23*e13 - (e33-eig[1])*e12)/bla
      eigv1[3] = (-e13 -e23*eigv1[2])/(e33-eig[1])
      aux = sqrt(1 + eigv1[2]^2 + eigv1[3]^2)
      eigv1[1] = 1/aux
      eigv1[2] = eigv1[2]/aux
      eigv1[3] = eigv1[3]/aux
    else
      bla = ((e11 - eig[1])*(e22 - eig[1]) - e12*e12)
      eigv1[3] = oneunit(T)
      eigv1[1] = (e23*e12 - (e22-eig[1])*e13)/bla
      eigv1[2] = (-e23 -e12*eigv1[1])/(e22-eig[1])
      aux = sqrt(1 + eigv1[2]^2 + eigv1[1]^2)
      eigv1[1] = eigv1[1]/aux
      eigv1[2] = eigv1[2]/aux
      eigv1[3] = 1/aux
    end
    bla = ((e22 - eig[2])*(e33 - eig[2]) - e23*e23)
    if bla != 0
      eigv2[1] = oneunit(T)
      eigv2[2] = (e23*e13 - (e33-eig[2])*e12)/bla
      eigv2[3] = (-e13 -e23*eigv2[2])/(e33-eig[2])
      aux = sqrt(1 + eigv2[2]^2 + eigv2[3]^2)
      eigv2[1] = 1/aux
      eigv2[2] = eigv2[2]/aux
      eigv2[3] = eigv2[3]/aux
    else
      bla = ((e11 - eig[2])*(e22 - eig[2]) - e12*e12)
      eigv2[3] = oneunit(T)
      eigv2[1] = (e23*e12 - (e22-eig[2])*e13)/bla
      eigv2[2] = (-e23 -e12*eigv2[1])/(e22-eig[2])
      aux = sqrt(1 + eigv2[2]^2 + eigv2[1]^2)
      eigv2[1] = eigv2[1]/aux
      eigv2[2] = eigv2[2]/aux
      eigv2[3] = 1.0/aux
    end

    eigv3[1] = eigv1[2]*eigv2[3] - (eigv1[3]*eigv2[2])
    eigv3[2] = eigv1[3]*eigv2[1] - (eigv1[1]*eigv2[3])
    eigv3[3] = eigv1[1]*eigv2[2] - (eigv1[2]*eigv2[1])

  end
  return 0
end

function eigen(e11::T,e22::T,e33::T,e12::T,e13::T,e23::T) where T<:AbstractFloat
  eig = MVector{3,T}()
  eigv1 = MVector{3,T}()
  eigv2 = MVector{3,T}()
  eigv3 = MVector{3,T}()
  eigen!(e11,e22,e33,e12,e13,e23,eig,eigv1,eigv2,eigv3)
  return SVector(eig), SVector(eigv1), SVector(eigv2), SVector(eigv3)
end

function eigen(e11::Real,e22::Real,e33::Real,e12::Real,e13::Real,e23::Real)
  return eigen(Float64(e11),Float64(e22),Float64(e33),Float64(e12),Float64(e13),Float64(e23))
end

function eigen(A::AbstractArray{<:Real,2})
  size(A) == (3,3) || throw(ArgumentError("Input must be a 3x3 Matrix"))
  (A[1,2] == A[2,1] && A[1,3] == A[3,1] && A[2,3] == A[3,2]) || throw(ArgumentError("Input must be symmetric matrix"))
  return eigen(A[1,1],A[2,2],A[3,3],A[1,2],A[1,3],A[2,3])
end

 function inphdecomp!(t11::T,t22::T,t33::T,t12::T,t13::T,t23::T,
                e11::T,e22::T,e33::T,e12::T,e13::T,e23::T,
                eig::MVector{3,T},eigv1::MVector{3,T},eigv2::MVector{3,T},eigv3::MVector{3,T}) where T<:AbstractFloat

  eigen!(e11,e22,e33,e12,e13,e23,eig, eigv1, eigv2, eigv3)

  m1 = muladd(t11, eigv1[1]*eigv1[1], muladd(t22, eigv1[2]*eigv1[2], muladd(t33, eigv1[3]*eigv1[3],
       2*muladd(t12, eigv1[1]*eigv1[2], muladd(t13, eigv1[1]*eigv1[3], t23*eigv1[2]*eigv1[3])))))
  m2 = muladd(t11, eigv2[1]*eigv2[1], muladd(t22, eigv2[2]*eigv2[2], muladd(t33, eigv2[3]*eigv2[3],
       2*muladd(t12, eigv2[1]*eigv2[2], muladd(t13, eigv2[1]*eigv2[3], t23*eigv2[2]*eigv2[3])))))
  m3 = t11 + t22 + t33 - m1 - m2

  lam2 = muladd(eig[1], eig[1], muladd(eig[2], eig[2], eig[3]*eig[3]))

  a = 1 - (3*(eig[1]^2)/lam2)
  b = 1 - (3*(eig[2]^2)/lam2)
  c = b*eig[1] - a*eig[2]

  alpha1 = (b*m1 - a*m2)/c
  alpha0 = (m1 - eig[1]*alpha1)/a
  alpha2 = -3*alpha0/lam2

  tm11 = muladd(m1, eigv1[1]*eigv1[1], muladd(m2, eigv2[1]*eigv2[1], m3*eigv3[1]*eigv3[1]))
  tm22 = muladd(m1, eigv1[2]*eigv1[2], muladd(m2, eigv2[2]*eigv2[2], m3*eigv3[2]*eigv3[2]))
  tm33 = muladd(m1, eigv1[3]*eigv1[3], muladd(m2, eigv2[3]*eigv2[3], m3*eigv3[3]*eigv3[3]))
  tm12 = muladd(m1, eigv1[1]*eigv1[2], muladd(m2, eigv2[1]*eigv2[2], m3*eigv3[1]*eigv3[2]))
  tm13 = muladd(m1, eigv1[1]*eigv1[3], muladd(m2, eigv2[1]*eigv2[3], m3*eigv3[1]*eigv3[3]))
  tm23 = muladd(m1, eigv1[2]*eigv1[3], muladd(m2, eigv2[2]*eigv2[3], m3*eigv3[2]*eigv3[3]))

  return tm11, tm22, tm33, tm12, tm13, tm23, alpha0, alpha1, alpha2
end

function inphdecomp(t11::T,t22::T,t33::T,t12::T,t13::T,t23::T,
                e11::T,e22::T,e33::T,e12::T,e13::T,e23::T) where T<:AbstractFloat

  eig = MVector{3,T}()
  eigv1 = MVector{3,T}()
  eigv2 = MVector{3,T}()
  eigv3 = MVector{3,T}()

  return inphdecomp!(t11,t22,t33,t12,t13,t23,e11,e22,e33,e12,e13,e23,eig,eigv1,eigv2,eigv3)
end

inphdecomp(var...) = inphdecomp((Float64(v) for v in var)...)

 function propdecomp(t11::T,t22::T,t33::T,t12::T,t13::T,t23::T,
                    e11::T,e22::T,e33::T,e12::T,e13::T,e23::T) where T<:AbstractFloat

  mode2 = (e11^2 + 2*e12^2 + 2*e13^2 + e22^2 + 2*e23^2 + e33^2)
  alphat = (e11*t11 + 2*e12*t12 + 2*e13*t13 + e22*t22 + 2*e23*t23 + e33*t33)/mode2

  m12 = alphat*e12
  m13 = alphat*e13
  m23 = alphat*e23
  m11 = alphat*e11
  m22 = alphat*e22
  m33 = alphat*e33

  return m11, m22, m33, m12, m13, m23, alphat
end

propdecomp(var...) = propdecomp((Float64(v) for v in var)...)

 function inphmodel(t11::Vector{T},t22::Vector{T},t33::Vector{T},t12::Vector{T},t13::Vector{T},t23::Vector{T}, #input
                   e11::Vector{T},e22::Vector{T},e33::Vector{T},e12::Vector{T},e13::Vector{T},e23::Vector{T}, #input
                   tm11::Vector{T},tm22::Vector{T},tm33::Vector{T},tm12::Vector{T},tm13::Vector{T},tm23::Vector{T}, #output
                   alpha0::Vector{T},alpha1::Vector{T},alpha2::Vector{T},ratio::Vector{T}) where T<:Real

  eigg  = [MVector{3,T}() for i in 1:Threads.nthreads()]
  eigv1 = [MVector{3,T}() for i in 1:Threads.nthreads()]
  eigv2 = [MVector{3,T}() for i in 1:Threads.nthreads()]
  eigv3 = [MVector{3,T}() for i in 1:Threads.nthreads()]

  Threads.@threads for i in 1:length(e11)
    j = Threads.threadid()
    @inbounds tm11[i], tm22[i], tm33[i], tm12[i], tm13[i], tm23[i], alpha0[i], alpha1[i], alpha2[i] = inphdecomp!(t11[i],t22[i],t33[i],t12[i],t13[i],t23[i],
                            e11[i],e22[i],e33[i],e12[i],e13[i],e23[i],
                            eigg[j],eigv1[j],eigv2[j],eigv3[j])

    @inbounds ratio[i] = (tm11[i]^2 + 2*tm12[i]^2 + 2*tm13[i]^2 + tm22[i]^2 + 2*tm23[i]^2 + tm33[i]^2)/(t11[i]^2 + 2*t12[i]^2 + 2*t13[i]^2 + t22[i]^2 + 2*t23[i]^2 + t33[i]^2)
  end
  return 0
end

 function propmodel(t11::A,t22::A,t33::A,t12::A,t13::A,t23::A, #input
                   e11::A,e22::A,e33::A,e12::A,e13::A,e23::A, #input
                   tm11::A,tm22::A,tm33::A,tm12::A,tm13::A,tm23::A, #output
                   alpha::A,ratio::A) where {A<:Vector{<:Real}}

  Threads.@threads for i in 1:length(e11)
    @inbounds tm11[i], tm22[i], tm33[i], tm12[i], tm13[i], tm23[i], alpha[i] = propdecomp(t11[i],t22[i],t33[i],t12[i],t13[i],t23[i],e11[i],e22[i],e33[i],e12[i],e13[i],e23[i])

    @inbounds ratio[i] = (tm11[i]^2 + 2*tm12[i]^2 + 2*tm13[i]^2 + tm22[i]^2 + 2*tm23[i]^2 + tm33[i]^2)/(t11[i]^2 + 2*t12[i]^2 + 2*t13[i]^2 + t22[i]^2 + 2*t23[i]^2 + t33[i]^2)
  end
  return 0
end

 function mixmodel(t11::A,t22::A,t33::A,t12::A,t13::A,t23::A, #input
                  e11::A,e22::A,e33::A,e12::A,e13::A,e23::A, #input
                  c11::A,c22::A,c33::A,c12::A,c13::A,c23::A, #input
                  tm11::A,tm22::A,tm33::A,tm12::A,tm13::A,tm23::A, #output
                  ratio::A) where {A<:Vector{<:Real}}

  Threads.@threads for i in 1:length(e11)
    @inbounds tm11[i] = e11[i] + c11[i]
    @inbounds tm22[i] = e22[i] + c22[i]
    @inbounds tm33[i] = e33[i] + c33[i]
    @inbounds tm12[i] = e12[i] + c12[i]
    @inbounds tm13[i] = e13[i] + c13[i]
    @inbounds tm23[i] = e23[i] + c23[i]

    @inbounds ratio[i] = (tm11[i]^2 + 2*tm12[i]^2 + 2*tm13[i]^2 + tm22[i]^2 + 2*tm23[i]^2 + tm33[i]^2)/(t11[i]^2 + 2*t12[i]^2 + 2*t13[i]^2 + t22[i]^2 + 2*t23[i]^2 + t33[i]^2)
  end
  return 0
end

end # module


function stepAdvection!(
    mb :: ModelBlock,
    Δt :: Float64,
)

    ev = mb.ev
    fi = mb.fi
    tmpfi = mb.tmpfi
    co = mb.co
    cfg = ev.config
    for x = 1:2

        _intmx  = view(tmpfi._INTMX_, :, x)
        _xflx_U = view(fi._Xflx_U_, :, x)
        _xflx_V = view(fi._Xflx_V_, :, x)
        _xflx_W = view(fi._Xflx_W_, :, x)

        calDiffAdv_QUICKEST!(
            _xflx_U,
            _xflx_V,
            _xflx_W,

            _intmx,

            fi._u,
            fi._v,
            fi._w,

            co.amo,
            reshape(co.mtx[:Ks_H_U], :),
            reshape(co.mtx[:Ks_H_V], :),
            0.0,
            Δt,

            co.wksp,
        )

        _ADVx_ = view(fi._ADVX_, :, x)

        _ADVx_[:] = - (
              co.amo.T_DIVx_U * _xflx_U
            + co.amo.T_DIVy_V * _xflx_V
            + co.amo.T_DIVz_W * _xflx_W
        )

        @. _intmx += Δt * _ADVx_

    end

end


f(theta)               = simloglik!(grad, info_matrix, tmpvars, model, theta, data, false)
g!(grad, theta)        = simloglik!(grad, info_matrix, tmpvars, model, theta, data, true)
h!(info_matrix, theta) = simloglik!(grad, info_matrix, tmpvars, model, theta, data, true)
fg! = g!
newton!(f, g, h, theta, x0)


using CompScienceMeshes, BEAST
using LinearAlgebra

T = tetmeshsphere(1.0,0.25)
X = nedelecc3d(T)
Γ = boundary(T)

X = raviartthomas(Γ)
@show numfunctions(X)

κ,  η  = 1.0, 1.0
κ′, η′ = 2.0κ, η/2.0

T  = Maxwell3D.singlelayer(wavenumber=κ)
T′ = Maxwell3D.singlelayer(wavenumber=κ′)
K  = Maxwell3D.doublelayer(wavenumber=κ)
K′ = Maxwell3D.doublelayer(wavenumber=κ′)

E = Maxwell3D.planewave(direction=ẑ, polarization=x̂, wavenumber=κ)
H = -1/(im*κ*η)*curl(E)

e, h = (n × E) × n, (n × H) × n

@hilbertspace j m
@hilbertspace k l

α, α′ = 1/η, 1/η′
pmchwt = @discretise(
    (η*T+η′*T′)[k,j] -      (K+K′)[k,m] +
         (K+K′)[l,j] + (α*T+α′*T′)[l,m] == -e[k] - h[l],
    j∈X, m∈X, k∈X, l∈X)

# Z = BEAST.sysmatrix(pmchwt)
# b = BEAST.rhs(pmchwt)
# M = BEAST.convert_to_dense(Z)
# u = M \ b
# error()
# W = BEAST.assemble_hide(pmchwt.equation.lhs, pmchwt.test_space_dict, pmchwt.trial_space_dict)

# using BEAST.BlockArrays
# using BEAST.IterativeSolvers

# u = PseudoBlockVector{ComplexF64}(undef, numfunctions.([X,X]))
# fill!(u,0)

# error()
# u, ch = IterativeSolvers.gmres!(u, Z, b, log=true,  maxiter=1000,
#     restart=1000, reltol=1e-5, verbose=true)

u = solve(pmchwt)

Θ, Φ = range(0.0,stop=2π,length=100), 0.0
ffpoints = [point(cos(ϕ)*sin(θ), sin(ϕ)*sin(θ), cos(θ)) for θ in Θ for ϕ in Φ]

# Don't forget the far field comprises two contributions
ffm = potential(MWFarField3D(κ*im), ffpoints, u[m], X)
ffj = potential(MWFarField3D(κ*im), ffpoints, u[j], X)
ff = -η*im*κ*ffj + im*κ*cross.(ffpoints, ffm)

using Plots
plot(xlabel="theta")
plot!(Θ,norm.(ff),label="far field",title="PMCHWT")

#import Plotly
#using LinearAlgebra
#fcrj, _ = facecurrents(u[j],X)
#fcrm, _ = facecurrents(u[m],X)
#Plotly.plot(patch(Γ, norm.(fcrj)))
#Plotly.plot(patch(Γ, norm.(fcrm)))

function nearfield(um,uj,Xm,Xj,κ,η,points,
    Einc=(x->point(0,0,0)),
    Hinc=(x->point(0,0,0)))

    K = BEAST.MWDoubleLayerField3D(wavenumber=κ)
    T = BEAST.MWSingleLayerField3D(wavenumber=κ)

    Em = potential(K, points, um, Xm)
    Ej = potential(T, points, uj, Xj)
    E = -Em + η * Ej + Einc.(points)

    Hm = potential(T, points, um, Xm)
    Hj = potential(K, points, uj, Xj)
    H = 1/η*Hm + Hj + Hinc.(points)

    return E, H
end

Z = range(-1,1,length=100)
Y = range(-1,1,length=100)
nfpoints = [point(0,y,z) for z in Z, y in Y]

import Base.Threads: @spawn
task1 = @spawn nearfield(u[m],u[j],X,X,κ,η,nfpoints,E,H)
task2 = @spawn nearfield(-u[m],-u[j],X,X,κ′,η′,nfpoints)

E_ex, H_ex = fetch(task1)
E_in, H_in = fetch(task2)

E_tot = E_in + E_ex
H_tot = H_in + H_ex

contour(real.(getindex.(E_tot,1)))
contour(real.(getindex.(H_tot,2)))

heatmap(Z, Y, real.(getindex.(E_tot,1)))
#heatmap(Z, Y, real.(getindex.(H_tot,2)))

#plot(real.(getindex.(E_tot[:,51],1)))
#plot!(real.(getindex.(H_tot[:,51],2)))

# Compare the far field and the field far
ffradius = 100.0
E_far, H_far = nearfield(u[m],u[j],X,X,κ,η, ffradius .* ffpoints)
nxE_far = cross.(ffpoints, E_far) * (4π*ffradius) / exp(-im*κ*ffradius)
Et_far = -cross.(ffpoints, nxE_far)

plot()
plot!(Θ, norm.(ff),label="far field")
scatter!(Θ, norm.(Et_far), label="field far")


@testset "gaussian" begin
    seed!(42)
    @testset "Likelihood" begin
        σ² = 1e-3
        l = GaussianLikelihood(σ²)
        f = 0.5
        y = 0.2
        @test AGP.noise(l) == σ²
        @test l(y, f) == pdf(Normal(f, sqrt(σ²)), y)
        @test loglikelihood(l, y, f) == logpdf(Normal(f, sqrt(σ²)), y)
        @test repr(l) == "Gaussian likelihood (σ² = $(AGP.noise(l)))"
        @test AGP.AugmentedKL(l, [], (;)) == 0.0
        @test AGP.n_latent(l) == 1
    end
    N, d = 20, 2
    k = SqExponentialKernel() ∘ ScaleTransform(10.0)
    σ = 0.1
    X, f = generate_f(N, d, k)
    y = f + σ * randn(N)
    floattypes = [Float64]
    @testset "GP" begin
        for floattype in floattypes
            model = GP(X, y, k; opt_noise=true, verbose=0)
            @test eltype(model) == floattype
            @test AGP.likelihood(model) isa GaussianLikelihood
            @test AGP.inference(model) isa Analytic
            @test AGP.getf(model) isa AGP.LatentGP
            @test AGP.n_latent(model) == 1
            L = AGP.objective(model, X, y)
            @test_nowarn train!(model, 10)
            @test L < AGP.objective(model, X, y)
            @test testconv(model, "Regression", X, f, y)
            @test all(proba_y(model, X)[2] .> 0)
        end
    end
    @testset "VGP" begin
        @test_throws ErrorException VGP(X, y, k, GaussianLikelihood(), AnalyticVI())
        @test_throws ErrorException VGP(X, y, k, GaussianLikelihood(), QuadratureVI())
        @test_throws ErrorException VGP(X, y, k, GaussianLikelihood(), MCIntegrationVI())
    end
    @testset "SVGP" begin
        Z = inducingpoints(KmeansAlg(M), X)
        for floattype in floattypes
            test_inference_SVGP(
                X,
                y,
                f,
                k,
                GaussianLikelihood(),
                GaussianLikelihood,
                floattype,
                Z,
                1,
                "Regression",
                AnalyticVI(),
                AnalyticSVI(10);
                valid=true,
            )
            @test_throws ErrorException SVGP(k, GaussianLikelihood(), QuadratureVI(), Z)
            @test_throws ErrorException SVGP(k, GaussianLikelihood(), MCIntegrationVI(), Z)
        end
    end
    @testset "MCGP" begin
        @test_throws ErrorException MCGP(X, y, k, GaussianLikelihood(), GibbsSampling())
        @test_throws ErrorException MCGP(X, y, k, GaussianLikelihood(), HMCSampling())
    end
end


using ColPack
using MatrixMarket
using SparseArrays
using SparseDiffTools
using LinearAlgebra

const filename = "/scratch/jac_case9241pegase_17036.mtx"

orderings = Vector{String}()

push!(orderings, "NATURAL")
push!(orderings, "LARGEST_FIRST")
push!(orderings, "DYNAMIC_LARGEST_FIRST")
push!(orderings, "DISTANCE_TWO_LARGEST_FIRST")
push!(orderings, "SMALLEST_LAST")
push!(orderings, "DISTANCE_TWO_SMALLEST_LAST")
push!(orderings, "INCIDENCE_DEGREE")
push!(orderings, "DISTANCE_TWO_INCIDENCE_DEGREE")
push!(orderings, "RANDOM")

verbose = 1
A = MatrixMarket.mmread(filename)

@time sparsediff_coloring = SparseDiffTools.matrix_colors(A)
ncolor = size(unique(sparsediff_coloring),1)
@time colpack_coloring = ColPackColoring("/scratch/jac_case9241pegase_17036.mtx", ColPack.d1, "RANDOM", verbose)

for ordering in orderings
    colpack_coloring = ColPackColoring("/scratch/jac_case9241pegase_17036.mtx", ColPack.d1, ordering, verbose)
end
colors = get_colors(colpack_coloring)
@show length(unique(colors))


using Test
using TestSetExtensions
using LinearAlgebra
using Qaintessent

##==----------------------------------------------------------------------------------------------------------------------

# adapted from https://github.com/FluxML/Zygote.jl/blob/master/test/gradcheck.jl
function ngradient(f, xs::AbstractArray...)
    grads = zero.(xs)
    for (x, Δ) in zip(xs, grads), i in 1:length(x)
        δ = sqrt(eps())
        tmp = x[i]
        x[i] = tmp - δ/2
        y1 = f(xs...)
        x[i] = tmp + δ/2
        y2 = f(xs...)
        x[i] = tmp
        Δ[i] = (y2-y1)/δ
        if eltype(x) <: Complex
            # derivative with respect to imaginary part
            x[i] = tmp - im*δ/2
            y1 = f(xs...)
            x[i] = tmp + im*δ/2
            y2 = f(xs...)
            x[i] = tmp
            Δ[i] += im*(y2-y1)/δ
        end
    end
    return grads
end

##==----------------------------------------------------------------------------------------------------------------------

@testset ExtendedTestSet "gate gradients" begin

    @testset "single qubit gates" begin

        # fictitious gradients of cost function with respect to quantum gate
        Δ = randn(ComplexF64, 2, 2)

        for g in [RxGate, RyGate, RzGate]
            f = (θ) -> 2*real(sum(Δ .* Qaintessent.sparse_matrix(g(θ[]))))
            θ = 2π*rand()
            ngrad = ngradient(f, [θ])
            dg = Qaintessent.backward(g(θ), conj(Δ))
            @test isapprox(dg.θ, ngrad[1], rtol=1e-5, atol=1e-5)
        end

        for g in [XGate, YGate, ZGate, HadamardGate, SGate, SdagGate, TGate, TdagGate]
            dg = Qaintessent.backward(g(), conj(Δ))
            @test dg ≈ g()
        end

        begin
            f = (ϕ) -> 2*real(sum(Δ .* Qaintessent.sparse_matrix(PhaseShiftGate(ϕ[]))))
            ϕ = 2π*rand()
            ngrad = ngradient(f, [ϕ])
            dg = Qaintessent.backward(PhaseShiftGate(ϕ), conj(Δ))
            @test isapprox(dg.ϕ, ngrad[1], rtol=1e-5, atol=1e-5)
        end

        begin
            f = (nθ) -> 2*real(sum(Δ .* Qaintessent.sparse_matrix(RotationGate(nθ))))
            nθ = randn(3)
            ngrad = ngradient(f, nθ)
            dg = Qaintessent.backward(RotationGate(nθ), conj(Δ))
            @test isapprox(dg.nθ, ngrad[1], rtol=1e-6)
        end

        begin
            f = (nθ) -> 2*real(sum(Δ .* Qaintessent.sparse_matrix(RotationGate(nθ))))
            # special case: zero vector
            nθ = zeros(3)
            ngrad = ngradient(f, nθ)
            dg = Qaintessent.backward(RotationGate(nθ), conj(Δ))
            @test isapprox(dg.nθ, ngrad[1], rtol=1e-6)
        end
    end

    @testset "two qubit gates" begin
        # fictitious gradients of cost function with respect to quantum gate
        Δ = randn(ComplexF64, 4, 4)
        @testset "entanglement gates" begin
            for g in [EntanglementXXGate, EntanglementYYGate, EntanglementZZGate]
                f(θ) = 2*real(sum(Δ .* Qaintessent.sparse_matrix(g(θ[]))))
                θ = 2π*rand()
                ngrad = ngradient(f, [θ])
                dg = Qaintessent.backward(g(θ), conj(Δ))
                @test isapprox(dg.θ, ngrad[1], rtol=1e-5, atol=1e-5)
            end
        end

        @testset "two qubit gates" begin
            for g in [SwapGate()]
                dg = Qaintessent.backward(g, conj(Δ))
                @test dg ≈ g
            end
        end

    end

    @testset "controlled gates" begin
        # fictitious gradients of cost function with respect to quantum gate
        Δ = randn(ComplexF64, 8, 8)

        for g in [RxGate, RyGate, RzGate]
            f(θ) = 2*real(sum(Δ .* Qaintessent.sparse_matrix(ControlledGate{g}(g(θ[]), 2))))
            θ = 2π*rand()
            ngrad = ngradient(f, [θ])
            dg = Qaintessent.backward(ControlledGate{g}(g(θ), 2), conj(Δ))
            @test isapprox(dg.U.θ, ngrad[1], rtol=1e-5, atol=1e-5)
        end

        for g in [EntanglementXXGate, EntanglementYYGate, EntanglementZZGate]
            f(θ) = 2*real(sum(Δ .* Qaintessent.sparse_matrix(ControlledGate{g}(g(θ[]), 1))))
            θ = 2π*rand()
            ngrad = ngradient(f, [θ])
            dg = Qaintessent.backward(ControlledGate{g}(g(θ), 2), conj(Δ))
            @test isapprox(dg.U.θ, ngrad[1], rtol=1e-5, atol=1e-5)
        end

        for g in [XGate, YGate, ZGate, HadamardGate, SGate, SdagGate, TGate, TdagGate]

            dg = Qaintessent.backward(ControlledGate{g}(g(), 2), conj(Δ))

            @test dg ≈ ControlledGate{g}(g(), 2)
        end
    end
end

##==----------------------------------------------------------------------------------------------------------------------

@testset ExtendedTestSet "circuit gradients" begin

    # construct parametrized circuit
    N = 4
    A = rand(ComplexF64, 2 ,2)
    U, R = qr(A)
    U = Array(U)
    i = rand(1:N)

    cgc(θ, ϕ, χ, κ, ωn) = CircuitGate[
        circuit_gate(3, HadamardGate()),
        circuit_gate(2, RzGate(θ), (1, 4)),
        circuit_gate(2, TGate(), 4),
        circuit_gate(2, 3, SwapGate()),
        circuit_gate(2, SdagGate(), 1),
        circuit_gate(3, PhaseShiftGate(ϕ)),
        circuit_gate(1, RyGate(χ)),
        circuit_gate(4, HadamardGate(), 2),
        circuit_gate(3, X, 1),
        circuit_gate(2, Z, 4),
        circuit_gate(4, Y, 1),
        circuit_gate(3, 1, EntanglementYYGate(κ)),
        circuit_gate(3, RotationGate(ωn)),
        circuit_gate(4, SGate()),
        circuit_gate(2, TdagGate(), 3),
        circuit_gate(i, MatrixGate(U)),
        circuit_gate(2, Z),
        circuit_gate(3, X),
        circuit_gate(2, Y, 4),
        circuit_gate(1, SGate()),
        circuit_gate(2, 3, SwapGate(), 4),
        circuit_gate(3, SdagGate()),
        circuit_gate(2, TGate()),
    ]

    # measurement operators
    meas(M) = MeasurementOperator.([Matrix{Float64}(I, 2^N, 2^N), Hermitian(M)], (Tuple(1:N),))

    # parameter values
    θ = 1.5π
    ϕ = 0.3
    χ = √2
    κ = exp(-0.4)
    n = randn(Float64, 3)
    n /= norm(n)
    ωn = 0.2π * n
    M = randn(ComplexF64, 2^N, 2^N)
    M = 0.5*(M + adjoint(M))

    c = Circuit{N}(cgc(θ, ϕ, χ, κ, ωn), meas(M))
    # input quantum state
    ψ = randn(ComplexF64, 2^N)
    ψl = deepcopy(ψ)
    # fictitious gradients of cost function with respect to circuit output
    Δ = [0.3, -1.2]

    dc = Qaintessent.gradients(c, ψ, Δ)[1]

    f(rθ, rϕ, rχ, rκ, ωn, M) = dot(Δ, apply(ψ, Circuit{N}(cgc(rθ[], rϕ[], rχ[], rκ[], ωn), meas(M))))
    # numeric gradients
    ngrad = ngradient(f, [θ], [ϕ], [χ], [κ], ωn, M)
    # symmetrize gradient with respect to measurement operator
    ngrad[end][:] = 0.5*(ngrad[end] + adjoint(ngrad[end]))

    @test all(isapprox.(ngrad,
        (dc.moments[1][2].gate.U.θ,
         dc.moments[4][2].gate.ϕ,
         dc.moments[5][1].gate.θ,
         dc.moments[8][1].gate.θ,
         dc.moments[9][1].gate.nθ,
         sparse_matrix(dc.meas[2])), rtol=1e-5, atol=1e-5))
end

##==----------------------------------------------------------------------------------------------------------------------

@testset ExtendedTestSet "circuit gradients with moments" begin

   # construct parametrized circuit
   N = 4
    cgc(θ, ϕ, χ, ωn) = Moment[
        Moment([
            circuit_gate(3, HadamardGate()),
            circuit_gate(2, RzGate(θ), (1, 4)),
        ]),
        Moment(circuit_gate(2, 3, SwapGate())),
        Moment(circuit_gate(3, PhaseShiftGate(ϕ))),
        Moment([
            circuit_gate(3, RotationGate(ωn)),
            circuit_gate(1, RyGate(χ)),
        ]),
    ]

   # measurement operators
   meas(M) = MeasurementOperator.([Matrix{Float64}(I, 2^N, 2^N), Hermitian(M)], (Tuple(1:N),))

   # parameter values
   θ = 1.5π
   ϕ = 0.3
   χ = √2
   n = randn(Float64, 3)
   n /= norm(n)
   ωn = 0.2π * n
   M = randn(ComplexF64, 2^N, 2^N)
   M = 0.5*(M + adjoint(M))

   c = Circuit{N}(cgc(θ, ϕ, χ, ωn), meas(M))

   # input quantum state
   ψ = randn(ComplexF64, 2^N)

   # fictitious gradients of cost function with respect to circuit output
   Δ = [0.3, -1.2]

   dc = Qaintessent.gradients(c, ψ, Δ)[1]

   f(rθ, rϕ, rχ, ωn, M) = dot(Δ, apply(ψ, Circuit{N}(cgc(rθ[], rϕ[], rχ[], ωn), meas(M))))
   # numeric gradients
   ngrad = ngradient(f, [θ], [ϕ], [χ], ωn, M)
   # symmetrize gradient with respect to measurement operator
   ngrad[end][:] = 0.5*(ngrad[end] + adjoint(ngrad[end]))

   @test all(isapprox.(ngrad,
       (dc.moments[1][2].gate.U.θ,
        dc.moments[3][1].gate.ϕ,
        dc.moments[4][2].gate.θ,
        dc.moments[4][1].gate.nθ,
        sparse_matrix(dc.meas[2])), rtol=1e-5, atol=1e-5))
end


struct PointMass{T<:Real} <: Prior
    atom::T
    PointMass{T}(atom::T) where{T<:Real} = (0 <= atom <= 1) ? new(atom) : error("atom must be in [0, 1]")
end
PointMass(atom::T) where{T<:Real} = PointMass{T}(atom)

bounds(prior::PointMass) = [prior.atom, prior.atom]

pdf(p::T, prior::PointMass{T}) where{T<:Real} = (p == prior.atom) ? Inf : 0

cdf(p::T, prior::PointMass{T}) where{T<:Real} = (p >= prior.atom) ? 1 : 0

function condition(prior::PointMass{T}; low::T = prior.atom, high::T = prior.atom) where {T<:Real}
    (low <= prior.atom <= high) ? (return prior) : error("conditioning only well-defined when probability atom is contained in intervsl")
end

update(prior::PointMass{T}, x::TI, n::TI) where {T<:Real,TI<:Integer} = !(0 <= x <= n) ? error("invalid x / n") : (return prior)

expectation(f::Function, prior::PointMass{T}) where {T<:Real} = f(prior.atom)

mean(prior::PointMass{T}) where {T<:Real} = prior.atom

string(prior::PointMass) = @sprintf "PointMass(%.2f)" prior.atom

quantile(prior::PointMass, prob::Real) = prior.atom


import SwitchOnSafety: isinfeasible, isfeasible, isdecided, usestep
export sdpbound_γ, sdpbound_gamma

function sdpbisection_alg(s::StateDepDiscreteSwitchedLinearSystem, optimizer, lo, hi; verbose=1, tol=1e-5, step=0.5)
    @assert hi > lo
    while hi - lo ≥ tol
        γ_mid = midpoint(hi, lo, step)
        status = sdplyap_feasible_prog(γ_mid, s, optimizer, verbose=verbose-1)
        if verbose ≥ 1
            println("$γ_mid")
        end
        if !isdecided(status, true)
            if usestep(lo, hi)
                step *= 2
                continue
            end
            γ_midlo = min(midpoint(lo, γ_mid, step), shift(γ_mid, -tol/2))
            statuslo = sdplyap_feasible_prog(γ_midlo, s, optimizer, verbose=verbose-1)
            if isdecided(statuslo, true)
                γ_mid = γ_midlo
                status = statuslo
            else
                γ_midhi = max(midpoint(γ_mid, hi, step), shift(γ_mid, tol/2))
                statushi = sdplyap_feasible_prog(γ_midhi, s, optimizer, verbose=verbose-1)
                if isdecided(statushi, true)
                    γ_mid = γ_midhi
                    status = statushi
                end
            end
        end
        if isinfeasible(status, true)
            lo = γ_mid
        end
        if isfeasible(status, true)
            hi = γ_mid
        end
    end
    return lo, hi
end

"""
    sdpbound_γ(s::StateDepDiscreteSwitchedLinearSystem)

Compute minimum value of ``\\gamma`` for which the SDP is feasible.

For keywords, see [`sosbound_γ`](@ref)
"""
function sdpbound_γ(s::StateDepDiscreteSwitchedLinearSystem; optimizer=optimizer_with_attributes(CSDP.Optimizer, MOI.Silent() => true), tol=1e-5, verbose=0, initstep=1.1)
    lo_init, hi_init = 0, initstep; hi_stat = (MOI.INFEASIBLE, MOI.INFEASIBLE_POINT, MOI.INFEASIBILITY_CERTIFICATE, "")
    while isinfeasible(hi_stat, true)
        hi_stat = sdplyap_feasible_prog(hi_init, s, optimizer, verbose=verbose-1)
        if isinfeasible(hi_stat, true)
            lo_init = copy(hi_init)
            hi_init += initstep
        end
        initstep *= 2
    end
    lo, hi = sdpbisection_alg(s, optimizer, lo_init, hi_init, verbose=verbose)
    return hi
end

"""
    sdpbound_gamma(s::StateDepDiscreteSwitchedLinearSystem; optimizer=nothing, tol=1e-5, verbose=0, initstep=1.1)

Alias for [`sdpbound_γ`](@ref).
"""
sdpbound_gamma(s::StateDepDiscreteSwitchedLinearSystem; optimizer=optimizer_with_attributes(CSDP.Optimizer, MOI.Silent() => true), tol=1e-5, verbose=0, initstep=1.1) = sdpbound_γ(s::StateDepDiscreteSwitchedLinearSystem; optimizer=optimizer, tol=tol, verbose=verbose, initstep=initstep)


function tracem(x::Array{T,2}) where T <: AbstractFloat

  # We require the number of rows to be greater than the number of columns, so that m is greater than one.

  trans = false

  if size(x,1) < size(x,2)
    x = x'
    trans = true
  end

  n = size(x,2)
  m = int(size(x,1)/n)

  y = zeros(m,1)  # We want this to be a 2-d array for subsequent matrix multiplication

  for i = 1:m

    y[i,1] = trace(x[(n*(i-1)+1):i*n,1:n])

  end

  if trans == true
    y = y'
  end

  return y

end

function permutation(n1::S,n2::S) where S <: Int

  nn = n1*n2
  p = zeros(S,nn,nn)

  i = 1
  j = 1
  for r = 1:nn

    if j > n2
      j = 1
      i += 1
    end
    c = i+n1*(j-1)
    p[r,c] = 1

    j += 1

  end

  return p

end

function convert_second_order(model::M) where M <: Lombardo_Sutherland_Soln

  ssh       = copy(model.ssh)
  hx        = copy(model.hx)
  hxx       = copy(model.hxx)
  ssg       = copy(model.ssg)
  gx        = copy(model.gx)
  gxx       = copy(model.gxx)
  sigma     = copy(model.sigma)
  grc       = copy(model.grc)
  soln_type = copy(model.soln_type)

  (n1,n2) = size(hx)
  (n3,n4) = size(gx)

  new_hxx = reshape(vech_to_vec(n2)*hxx[1,:]',n2,n2)
  for i = 2:n1

    new_hxx = [new_hxx; reshape(vech_to_vec(n2)*hxx[i,:]',n2,n2)]

  end

  (n5,n6) = size(new_hxx)

  j = 0
  for i = 1:n5

    j += 1
    new_hxx[i,j] = 2*new_hxx[i,j]
    if j == n6
      j = 0
    end

  end

  new_hxx = new_hxx/2

  new_gxx = reshape(vech_to_vec(n4)*gxx[1,:]',n4,n4)
  for i = 2:n3

    new_gxx = [new_gxx; reshape(vech_to_vec(n4)*gxx[i,:]',n4,n4)]

  end

  (n7,n8) = size(new_gxx)

  j = 0
  for i = 1:n7

    j += 1
    new_gxx[i,j] = 2*new_gxx[i,j]
    if j == n8
      j = 0
    end

  end

  new_gxx = new_gxx/2

  soln = Gomme_Klein_Soln(ssh,hx,new_hxx,ssg,gx,new_gxx,sigma,grc,soln_type)

  return soln

end


@testset "math" begin
   @testset "$f" for f in (cor,cov)
        A = rand(3, 5)
        Aindices = IArray(A, ([:one, :two, :three], 2:6))

       @testset "matrix input, matrix result" begin
            findices = f(Aindices, dims=1)
            fbase = f(A, dims=1)

            @test findices == fbase
            findices = keys.(axes(findices))

            @test findices == (2:6, 2:6)

            findices = f(Aindices, dims=2)
            fbase = f(A, dims=2)

            @test findices == fbase
            findices = keys.(axes(findices))

            @test findices == ([:one, :two, :three], [:one, :two, :three])
        end
        @testset "vector input, scalar result" begin
            A = rand(4)
            Aindices = IArray(A, 2:5)

            findices = f(Aindices)
            fbase = f(A)

            @test f(Aindices) isa Number
            @test findices == fbase
        end
    end
end

@testset "+" begin
    a = IArray(ones(3))
    @test +(a) == ones(3)
    @test +(a, a) == 2ones(3)
    @test +(a, a, a) == 3ones(3)
    @test +(a, parent(a)) == 2ones(3)
    @test +(parent(a), a) == 2ones(3)

        #=
        casts = (NamedDimsArray{(:foo, :bar)}, identity)
        for (T1, T2, T3, T4) in Iterators.product(casts, casts, casts, casts)
            all(isequal(identity), (T1, T2, T3, T4)) && continue
            total = T1(ones(3, 6)) + T2(2ones(3, 6)) + T3(3ones(3, 6)) + T4(4ones(3, 6))
            @test total == 10ones(3, 6)
            @test names(total) == (:foo, :bar)
        end
        =#
end

@testset "-" begin
    a = IArray(ones(3))
    @test -(a) == -ones(3)
    @test -(a, a) == -(parent(a), parent(a))
    @test -(a, a, a) == parent(a) - parent(a) - parent(a)
    @test -(a, parent(a)) == -(parent(a), parent(a))
    @test -(parent(a), a) == -(parent(a), parent(a))
end


function test_sampler_cube(sampler, n_samples; opts=Dict(), add_slice=nothing)
	""" Creates a 3-dimensional unit cube and uses the given sampler
		to sample from it the specified number of times. Then plots the
		sampled distribution
	"""
	A = Array{Float64}(undef, 0, 3)
	b = Array{Float64}(undef, 0, 1)
	p0 = [1.0, 0.5, 0.5]
	v0 = [1.0, 0.0, 0.0]
	p1 = [0, 0.5, 0.5]
	v1 = [-1.0, 0.0, 0.0]
	for (p, v) in zip((p0,p1), (v0,v1))
	for i in 0:2
		p_ = circshift(p, i)
		v_ = circshift(v, i)
		A = vcat(A, v_')
		b = vcat(b, v_' * p_)
	end
	end

	if !isnothing(add_slice)
		v_ = add_slice[1]
		p_ = add_slice[2]
		A = vcat(A, v_')
		b = vcat(b, v_' * p_)
	end

	samples = Array{Float64}(undef, 0, 3)

	x0, r = OPFLearn.chebyshev_center(A, b)

	k = 0
	while k < n_samples
		sample = sampler(A, b, x0'; opts...)'
		samples = vcat(samples, sample)
		k = k + size(sample, 1)
	end
	return samples
end

"""
using PyPlot
const plt = PyPlot
function plot_samples_series(samples, loads=(1, 2, 3); sampler="", case="")

	num_samples = size(samples,1)

    fig = plt.figure()
    ax = fig.add_subplot(projection="3d")

    # Plots with colors indicating sample number
    p = ax.scatter(samples[:, loads[1]],
                   samples[:, loads[2]],
                   samples[:, loads[3]],
                   c=collect(1: size(samples,1)), cmap="plasma")
    fig.colorbar(p, label="Sample Number")
    ax.set_xlabel("Load " * string(loads[1]))
    ax.set_ylabel("Load " * string(loads[2]))
    ax.set_zlabel("Load " * string(loads[3]))
    plt.title(case * " load powers, " * string(num_samples) * " samples, " * sampler * " sampler")
    plt.tight_layout()
    plt.show()
end

function run_test_sampler(num_samples, add_slice=[[1.0,1.0,1.0],[0.5,0.5,0.5]])
	opts = Dict((:method=>"achr"))
	samples = test_sampler_cube(sample_polytope_cprnd, num_samples, opts=opts, add_slice=add_slice)
	plot_samples_series(samples, sampler="cprnd", case="Unit Cube")
end
"""

@testset "sampler uniformity check" begin
    @testset "cprnd unit cube" begin
		num_samples = 10000  # 10,000 samples should take approximately 30 seconds
		opts = Dict((:method=>"achr"))
		samples = test_sampler_cube(sample_polytope_cprnd, num_samples, opts=opts, add_slice=nothing)
		@test size(samples, 1) == num_samples
		@test size(samples, 2) == 3
		@test isapprox(maximum(samples[:,1]), 1; atol = 1e-1)
		@test isapprox(maximum(samples[:,2]), 1; atol = 1e-1)
		@test isapprox(maximum(samples[:,3]), 1; atol = 1e-1)
		test_points = ((1.0, 1.0, 1.0),
					   (0.5, 0.5, 0.5),
					   (0.0, 0.0, 0.0),
					   [circshift([1.0, 0.0, 0.0],i) for i in 1:3]...,
					   [circshift([1.0, 1.0, 0.0],i) for i in 1:3]...,
					   [circshift([1.0, 0.5, 0.5],i) for i in 1:3]...,
					   [circshift([0.0, 0.5, 0.5],i) for i in 1:3]...,
					   )

		for p in test_points  # Test that at least one point sampled is near a test point
		@test isapprox(minimum([sum((sample .- p).^2) for sample in eachrow(samples)]), 0; atol = 5e-2)
		end
    end

    @testset "cprnd unit cube with slice" begin
		num_samples = 10000  # 10,000 samples should take approximately 30 seconds
		opts = Dict((:method=>"achr"))
		normal = [1.0, 1.0, 1.0]
		point = [0.5, 0.5, 0.0]
		slice = [normal, point]
		samples = test_sampler_cube(sample_polytope_cprnd, num_samples, opts=opts, add_slice=slice)
		@test size(samples, 1) == num_samples
		@test size(samples, 2) == 3
		@test isapprox(maximum(samples[:,1]), 1; atol = 1e-1)
		@test isapprox(maximum(samples[:,2]), 1; atol = 1e-1)
		@test isapprox(maximum(samples[:,3]), 1; atol = 1e-1)
		inside_test_points = ((0.0, 0.0, 0.0),
							  (0.5, 0.5, 0.0),
							  (0.25, 0.25, 0.25),
							  [circshift([1.0, 0.0, 0.0],i) for i in 1:3]...,
						      )

		outside_test_points = ((1.0, 1.0, 1.0),
							   (0.5, 0.5, 0.5),
							   [circshift([1.0, 1.0, 0.0],i) for i in 1:3]...,
							   )
		for p in inside_test_points  # Test that at least one point sampled is near a test point
		@test isapprox(minimum([sum((sample .- p).^2) for sample in eachrow(samples)]), 0; atol = 5e-2)
		end
		for p in outside_test_points  # Test that there are no points sampled is near a test point
		@test minimum([sum((sample .- p).^2) for sample in eachrow(samples)]) > 1e-2
		end
    end
end


using CuArrays

# TODO: custom kernel for activations on GPU

function sigmoid(x::Float32)::Float32
    return 1 / (1 + CuArrays.exp(-x))
end

function sigmoid_back(x::Float32)::Float32
    sig = sigmoid(x)
    return sig*(1-sig)
end

function relu(x::Float32)::Float32
    if x>0
        return x
    else
        return 0
    end
end

function relu_back(x::Float32)::Float32
    if x>0
        return 1
    else
        return 0
    end
end


using Transits
using Limbdark
using LaTeXStrings
using BenchmarkTools
using DataFrames
using Statistics
using StatsPlots, ColorSchemes
using ProgressLogging

N_pts = [100, 316, 1000, 3160, 10000, 31600, 100000]
N_us = [1 2 3 5 8 13 21 34 55 90 144]

function benchmark(N_u, N_pts; r=0.1)
    u = ones(N_u) ./ N_u
    b = [sqrt(((i - N_pts/2) * 2/N_pts * (1 + 2 * r))^2) for i in 1:N_pts]

    ld = PolynomialLimbDark(u)
    output = similar(b)
    bench = @benchmark begin
        @inbounds for i in eachindex($b)
            $output[i] = compute($ld, $b[i], $r)
        end
    end
    return median(bench).time * 1e-9
end
vals = zeros(length(N_pts), length(N_us))
@progress for (j, Nu) in enumerate(N_us)
    @progress for (i, N) in enumerate(N_pts)
        vals[i, j] = benchmark(Nu, N)
    end
end

labels = reduce(hcat, "uₙ: $N" for N in N_us)

plot(
    N_pts,
    vals,
    scale=:log10,
    ls=:solid,
    m=:o,
    palette=palette(:plasma, length(N_us) + 2),
    msw=0,
    lab=labels,
    xlabel="number of points",
    ylabel="timing [s]",
    leg=:topleft
)

function benchmark_limbdark(N_u, N_pts; r=0.1)
    u = ones(N_u) ./ N_u
    b = [sqrt(((i - N_pts/2) * 2/N_pts * (1 + 2 * r))^2) for i in 1:N_pts]

    trans = Limbdark.transit_init(r, b[1], u, true)
    output = similar(b)
    bench = @benchmark begin
        @inbounds for i in eachindex($b)
            $trans.b = $b[i]
            $output[i] = Limbdark.transit_poly!($trans)
        end
    end
    return median(bench).time * 1e-9
end
vals_limbdark = zeros(length(N_pts), length(N_us))
@progress for (j, Nu) in enumerate(N_us)
    @progress for (i, N) in enumerate(N_pts)
        vals_limbdark[i, j] = benchmark_limbdark(Nu, N)
    end
end

plot(
    N_pts,
    vals_limbdark ./ vals,
    xscale=:log10,
    ls=:solid,
    m=:o,
    msw=0,
    palette=palette(:plasma, length(N_us) + 2),
    lab=labels,
    xlabel="number of points",
    ylabel="relative speedup (times faster) - Limbdark.jl"
)

tmed = median(vals ./ vals[:, 1], dims=1)
plot(
    N_us |> vec,
    tmed |> vec,
    scale=:log10,
    ls=:solid,
    m=:o,
    msw=0,
    lab="Transits.jl",
    xlabel="number of limb-darkening coefficients",
    ylabel="relative timing",
    ylims=(0.8, 10),
    leg=:topleft
)
plot!(
    N_us |> vec,
    median(vals_limbdark ./ vals_limbdark[:, 1], dims=1) |> vec,
    lab="Limbdark.jl",
    c=2,
    ls=:solid,
    m=:o,
    msw=0
)
plot!(N_us |> vec, n -> n^(0.2), ls=:dash, c=:gray, lab=L"N^{0.2}")
plot!(N_us |> vec, n -> tmed[end] * n / N_us[1, end], ls=:dash, c=:gray, lab=L"N^1")


using SharedArrays
using JLD
include("../examples/lorenz63.jl")
function ergodic_average(rho)
	s = [10., rho, 8/3]
	n = 100000000000
	n_runup = 5000
	z = 0.
	u = rand(3)
	u .= lorenz63(reshape(u,3,1),
			s, n_runup)[end,:,1]
	for i = 1:n
		u .= next(u, s)
		z += u[3]/n
		if i==n
		  @time u .= next(u,s)
		end
	end
	return z-rho
end
function compute_erg_avg()
	n_p = 250
	rho = LinRange(28.0, 28.25, n_p)
	z_rho = SharedArray{Float64, 1}(n_p)
	t = @distributed for i = 1:n_p
		@show i
		z_rho[i] = ergodic_average(rho[i])
	end
	wait(t)
	@show z_rho
	save("../data/erg_avg/Jar1_res.jld", "z_rho",
	     z_rho, "rho", rho)

end


using Base.Test

import FracFin

# write your own tests here
# @test 1 == 2

# grid = 1:10^6

# fBM = FracFin.FractionalBrownianMotion(0.7)
# # sampler = FracFin.CholeskySampler(fBM, grid)
# sampler = FracFin.CircSampler(fBM, grid)
# # sampler = FracFin.CircSampler{FracFin.FractionalBrownianMotion, UnitRange{Int64}}(fBM, grid)
# X = rand(sampler)

# # fGN = FracFin.FractionalGaussianNoise(0.7)
# # sampler = FracFin.CircSampler(fGN, grid)
# # # sampler = CholeskySampler{FractionalGaussianNoise, UnitRange{Int64}}(fGN, grid)
# # X = FracFin.rand(sampler)

# # print(X)


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

@recipe function f(lpars::LocalAnisotropy, D::SpatialData, axis=:Y)
  @assert ndims(lpars) == 2 "plot only available for 2D local anisotropies"
  quats = lpars.rotation
  u = [quat_to_dcm(quats[x])[1,1] for x in 1:length(quats)]
  v = [quat_to_dcm(quats[x])[1,2] for x in 1:length(quats)]
  x = [coordinates(centroid(D,x))[1] for x in 1:nelements(D)]
  y = [coordinates(centroid(D,x))[2] for x in 1:nelements(D)]
  c = lpars.magnitude[iaxis(axis),:]

  seriestype --> :quiver
  quiver --> (u,v)
  line_z --> repeat(c, inner=4)
  if maximum(c) <= 1
    seriescolor --> :redsblues
  else
    seriescolor --> :bluesreds
  end

  x,y
end


module PackmolInputCreator

using PDBTools

export find_x
export convert_c
export write_input
export interpolate

# conversion factor from mL/mol to Å^3/molecule
const CMV = 1e24 / 6.02e23

# conversion factor from mol/L to molecules/Å^3
const CMC = 6.02e23 / 1e27

"""

Interpolate a value from a (x,y) data matrix

"""
function interpolate(x,ρ)
  i = findfirst(d -> d > x, ρ[:,1])
  dρdx = (ρ[i,2]-ρ[i-1,2])/(ρ[i,1]-ρ[i-1,1])
  d = ρ[i-1,2] + dρdx * (x - ρ[i-1,1])
  return d
end

"""

Find by bisection which is the molar fraction that is consistent
with a desired volume fraction (%), given a data table of densities
as a function of the molar fractions.

"""
function find_x(
  vv, cossolvent_mass, density_pure, densities::Matrix;
  tol = 1e-5, maxit=1000
)

  xl = densities[1,1]
  xr = densities[end,1]
  if xl < xr
    increasing = true
  else
    increasing = false
  end
  x = (xr-xl)/2
  xnew = convert_c(
    vv,"%vv" => "x", density=interpolate(x,densities),
    molar_mass=cossolvent_mass, density_pure=density_pure
  )
  it = 0
  while abs(xnew-x) > tol
    if xnew > x
      increasing ? xl = x : xr = x
    else
      increasing ? xr = x : xl = x
    end
    x = (xr+xl)/2
    xnew = convert_c(
      vv,"%vv" => "x", density=interpolate(x,densities),
      molar_mass=cossolvent_mass, density_pure=density_pure
    )
    it += 1
    it > maxit && error("Maximum number of iterations achieved.")
  end
  return x
end

"""

Convert concentrations.

"""
function convert_c(
  cin,units;
  density=nothing, # solution
  density_pure=nothing, # cossolvent
  molar_mass=18., # cossolvent
  molar_mass_water=18.,
  density_water=1.0
)

  # Why not allow this
  units[1] == units[2] && return cin

  ρ = density
  ρc = density_pure
  ρw = density_water
  Mw = molar_mass_water
  Mc = molar_mass

  # nc and nw are the molar concentrations

  if units[1] == "%vv"
    if ρc == nothing
      error("Density of pure solvent is required to convert from %vv.")
    end
    vv = cin/100
    nc = ( ρc*vv / Mc )
    if units[2] == "x"
      if ρ == nothing
        error("Density of solution is required to convert to molar fraction.")
      end
      nw = ( ρ - nc*Mc ) / Mw
      return nc/(nc + nw)
    end
    if units[2] == "mol/L"
      return 1000*nc
    end
  end

  if units[1] == "x"
    if ρ == nothing
      error("Density of solution is required to convert from molar fraction.")
    end
    x = cin
    nc = ρ / ( Mc + Mw*(1-x)/x )
    if units[2] == "mol/L"
      return 1000*nc
    end
    if units[2] == "%vv"
      if ρc == nothing
        error("Density of pure solvent is required to convert to %vv.")
      end
      nw = nc*(1-x)/x
      Vc = nc * Mc / ρc
      vv = 100*Vc
      return vv
    end
  end

  if units[1] == "mol/L"
    if ρ == nothing
      error("Density of solution is required to convert from molarity.")
    end
    nc = cin/1000
    nw = ( ρ - nc*Mc ) / Mw
    if units[2] == "x"
      return nc/(nc + nw)
    end
    if units[2] == "%vv"
      if ρc == nothing
        error("Density of pure solvent is required to convert to %vv.")
        return nothing
      end
      Vc = nc * Mc / ρc
      vv = 100*Vc
      return vv
    end
  end

end

"""

Function that generates an input file for Packmol. By default, the concentrations is given in mol/L, but it can also be given in molar fraction "x" or volume percentage "%vv", using `cunit="x"` or `cunit="%vv"`.

"""
function write_input(
  pdbfile::String, solvent_file::String, concentration::Real, box_side::Real;
  water_file="tip4p2005.pdb",
  density=1.0,
  density_pure_solvent=nothing,
  cunit="mol/L",
  packmol_input="box.inp",
  packmol_output="system.pdb"

)

  protein = readPDB(pdbfile)
  cossolvent = readPDB(solvent_file)

  # molar masses (g/mol)
  Mp = mass(protein)
  Mc = mass(cossolvent)
  Mw = 18.01528

  # aliases for clearer formulas
  ρ = density
  ρc = density_pure_solvent

  # Convert concentration to mol/L
  cc_mol = convert_c(
    concentration, cunit => "mol/L",
    density=ρ,
    density_pure=ρc,
    molar_mass=Mc
  )
  c_vv = convert_c(
    concentration, cunit => "%vv",
    density=ρ,
    density_pure=ρc,
    molar_mass=Mc
  )
  c_x = convert_c(
    concentration, cunit => "x",
    density=ρ,
    density_pure=ρc,
    molar_mass=Mc
  )

  # Convert cossolvent concentration in molecules/Å³
  cc = CMC*cc_mol

  # Box volume (Å³)
  vbox = box_side^3

  # Solution volume (vbox - vprotein)
  vs = vbox - CMV*Mp/ρ

  # number of cossolvent molecules
  nc = round(Int,cc*vs)

  #number of water molecules
  nw = round(Int,(ρ*vs/CMV - nc*Mc)/Mw)

  # Final density of the solution
  ρ = CMV*(Mc*nc + Mw*nw)/vs

  # Final cossolvent concentration (mol/L)
  cc_f = 1000*(nc/vs)*CMV

  # Final water concentration (mol/L)
  cw_f = 1000*(nw/vs)*CMV

  # Final recovered concentration in %vv
  vv = 100*CMV*(nc*Mc/ρc)/vs

  println(
    """
    Summary:
    ========
    Target concentration = $cc_mol mol/L
                         = $c_vv %vv
                         = $c_x x
                         = $cc molecules/Å³

    Box volume = $vbox Å³
    Solution volume = $vs Å³

    Density = $ρ g/mL
    Protein molar mass = $Mp g/mol
    Cossolvent molar mass = $Mc g/mol

    Number of cossolvent molecules = $nc molecules
    Number of water molecules = $nw molecules

    Final cossolvent concentration = $cc_f mol/L
    Final water concentration = $cw_f mol/L
                              = $(CMC*cw_f) molecules/Å³
    Final solvent density = $ρ g/mL
    Final %vv concentration = $vv %
    Final molar fraction = $(nc/(nc+nw))
    """)

  l = box_side/2
  open(packmol_input,"w") do io
    println(io,
      """
      tolerance 2.0
      output $packmol_output
      add_box_sides 1.0
      filetype pdb
      seed -1

      structure $pdbfile
        number 1
        center
        fixed 0. 0. 0. 0. 0. 0.
      end structure

      structure $water_file
        number $nw
        inside box -$l -$l -$l $l $l $l
      end structure
      """)
    if nc > 0
      println(io,
        """
        structure $solvent_file
          number $nc
          inside box -$l -$l -$l $l $l $l
        end structure
        """)
   end
 end
 println("Wrote file: $packmol_input")

end # function write_input

end # module


# exports
export create, Cuboid, show

mutable struct Cuboid <: Shape
    bounds::Array{Float64, 2}
end

function Base.show(io::IO, x::Cuboid)
    println(io, "Cuboid")
    for i in 1:size(x.bounds)[1]
        println(io, "Dimension $i: Low $(x.bounds[i, 1]) High $(x.bounds[i, 2])")
    end
end

function create(c::Cuboid; resolution=nothing, rand_=0.0, type::Int64=1)
    if isa(resolution, Nothing)
        N = [10 for i in 1:size(c.bounds)[1]]
    else
        N = Int64.(round.((c.bounds[:, 2] - c.bounds[:, 1])/resolution))
    end
    lattice = (c.bounds[:, 2] - c.bounds[:, 1]) ./ N
    mesh = zeros(3, prod(N))
    a = 1
    R = CartesianIndices(Tuple(N))
    for I in R
        I_ = Tuple(I)
        for i in eachindex(I_)
            mesh[i, a] = c.bounds[i, 1] - lattice[i]/2 + (I_[i] + rand_*randn())*lattice[i]
        end
        a += 1
    end
    return Dict(
        :x => mesh,
        :v => zeros(size(mesh)),
        :y => copy(mesh),
        :volume => ones(prod(N))*prod(lattice),
        :type => type*ones(Int64, prod(N))
        )
end


using Cbc
using Test
using LinearAlgebra

# using Gurobi
# solver = GurobiSolver(OutputFlag=0)

# using CPLEX
# solver = CplexSolver(CPX_PARAM_SCRIND=0)

import JuMP
import MathOptInterface
const MOI = MathOptInterface

using PiecewiseLinearOpt

# TODO: Add :SOS2 to this list, once Cbc supports it
methods_1D = (:CC,:MC,:Logarithmic,:LogarithmicIB,:ZigZag,:ZigZagInteger,:GeneralizedCelaya,:SymmetricCelaya,:Incremental,:DisaggLogarithmic)
# TODO: Add :SOS2 to this list, once Cbc supports it
# TODO: Add :MC to this list, Cbc (but not Gurobi) gives a different answer below, only for :MC (maybe a bug in Cbc?)
methods_2D = (:CC,:Logarithmic,:LogarithmicIB,:ZigZag,:ZigZagInteger,:GeneralizedCelaya,:SymmetricCelaya,:DisaggLogarithmic)
patterns_2D = (:Upper,:Lower,:BestFit,:UnionJack,:K1,:Random) # :OptimalTriangleSelection and :Stencil not supported currently

println("\nunivariate tests")
@testset "1D: $method" for method in methods_1D
    model = JuMP.Model(JuMP.with_optimizer(Cbc.Optimizer))
    JuMP.@variable(model, x)
    z = piecewiselinear(model, x, range(1,stop=2π, length=8), sin, method=method)
    JuMP.@objective(model, Max, z)

    JuMP.optimize!(model)

    @test JuMP.termination_status(model) == MOI.OPTIMAL
    @test JuMP.value(x) ≈ 1.75474 rtol=1e-4
    @test JuMP.value(z) ≈ 0.98313 rtol=1e-4

    JuMP.@constraint(model, x ≤ 1.5z)

    JuMP.optimize!(model)

    @test JuMP.termination_status(model) == MOI.OPTIMAL
    @test JuMP.value(x) ≈ 1.36495 rtol=1e-4
    @test JuMP.value(z) ≈ 0.90997 rtol=1e-4
    @test JuMP.objective_value(model) ≈ 0.90997 rtol=1e-4
    @test JuMP.objective_value(model) ≈ JuMP.value(z) rtol=1e-4
end

println("\nbivariate tests")
@testset "2D: $method, $pattern" for method in methods_2D, pattern in patterns_2D
    model = JuMP.Model(JuMP.with_optimizer(Cbc.Optimizer))
    JuMP.@variable(model, x[1:2])
    d = range(0,stop=1,length=8)
    f = (x1,x2) -> 2*(x1-1/3)^2 + 3*(x2-4/7)^4
    z = piecewiselinear(model, x[1], x[2], BivariatePWLFunction(d, d, f, pattern=pattern), method=method)
    JuMP.@objective(model, Min, z)

    JuMP.optimize!(model)

    @test JuMP.termination_status(model) == MOI.OPTIMAL
    @test JuMP.value(x[1]) ≈ 0.285714 rtol=1e-4
    @test JuMP.value(x[2]) ≈ 0.571429 rtol=1e-4
    @test JuMP.value(z) ≈ 0.004535 rtol=1e-3
    @test JuMP.objective_value(model) ≈ 0.004535 rtol=1e-3
    @test JuMP.objective_value(model) ≈ JuMP.value(z) rtol=1e-3

    JuMP.@constraint(model, x[1] ≥ 0.6)

    JuMP.optimize!(model)

    @test JuMP.termination_status(model) == MOI.OPTIMAL
    @test JuMP.value(x[1]) ≈ 0.6 rtol=1e-4
    @test JuMP.value(x[2]) ≈ 0.571428 rtol=1e-4
    @test JuMP.value(z) ≈ 0.148753 rtol=1e-4
    @test JuMP.objective_value(model) ≈ 0.148753 rtol=1e-3
    @test JuMP.objective_value(model) ≈ JuMP.value(z) rtol=1e-3
end
#
# println("\nbivariate optimal IB scheme tests")
# @testset "2D: optimal IB, UnionJack" begin
#     model = JuMP.Model(JuMP.with_optimizer(Cbc.Optimizer))
#     JuMP.@variable(model, x)
#     JuMP.@variable(model, y)
#     d = range(0,stop=1,length=3)
#     f = (x,y) -> 2*(x-1/3)^2 + 3*(y-4/7)^4
#     z = piecewiselinear(model, x, y, BivariatePWLFunction(d, d, f, pattern=:UnionJack), method=:OptimalIB, subsolver=solver)
#     JuMP.@objective(model, Min, z)
#
#     JuMP.optimize!(model)
#
#     @test JuMP.termination_status(model) == MOI.OPTIMAL
#     @test JuMP.value(x) ≈ 0.5 rtol=1e-4
#     @test JuMP.value(y) ≈ 0.5 rtol=1e-4
#     @test JuMP.value(z) ≈ 0.055634 rtol=1e-3
#     @test getobjectivevalue(model) ≈ 0.055634 rtol=1e-3
#     @test getobjectivevalue(model) ≈ JuMP.value(z) rtol=1e-3
#
#     JuMP.@constraint(model, x ≥ 0.6)
#
#     JuMP.optimize!(model)
#
#     @test JuMP.termination_status(model) == MOI.OPTIMAL
#     @test JuMP.value(x) ≈ 0.6 rtol=1e-4
#     @test JuMP.value(y) ≈ 0.5 rtol=1e-4
#     @test JuMP.value(z) ≈ 0.222300 rtol=1e-3
#     @test JuMP.objective_value(model) ≈ 0.222300 rtol=1e-3
#     @test JuMP.objective_value(model) ≈ JuMP.value(z) rtol=1e-3
# end


# output is pasted to README
using RandomSubSamplers, Random
Random.seed!(3);
v = collect(zip(rand(1:100, 1000), rand('a':'z', 1000)));
w = collect(random_subsample(first, 0.4, v)); # keep approx 40% of integer keys
length(w)
k = unique(first.(w));
length(k)                         # ≈ 40% (small-sample random noise)
filter(t -> first(t) ∈ k, v) == w # equivalent to selection, but lazy


# ---
# title: 508. Most Frequent Subtree Sum
# id: problem508
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Hash Table, Tree
# link: <https://leetcode.com/problems/most-frequent-subtree-sum/description/>
# hidden: true
# ---
#
# Given the root of a tree, you are asked to find the most frequent subtree sum.
# The subtree sum of a node is defined as the sum of all the node values formed
# by the subtree rooted at that node (including the node itself). So what is the
# most frequent subtree sum value? If there is a tie, return all the values with
# the highest frequency in any order.
#
# **Examples 1**
# Input:
#
#
#
#       5
#      /  \
#     2   -3
#
#
# return [2, -3, 4], since all the values happen only once, return all of them
# in any order.
#
# **Examples 2**
# Input:
#
#
#
#       5
#      /  \
#     2   -5
#
#
# return [2], since 2 happens twice, however -5 only occur once.
#
# **Note:** You may assume the sum of values in any subtree is in the range of
# 32-bit signed integer.
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


using Base.Test
import WaveletScattering: Log1P

x = [0.0, 1.0, 2.0, 3.0]
data_in = expm1.(x)
data_out = zeros(data_in)
ρ = Log1P(1.0)
Base.map!(ρ, data_out, data_in)
@test_approx_eq x data_out


immutable Laplace <: ContinuousUnivariateDistribution
    μ::Float64
    β::Float64

    function Laplace(μ::Real, β::Real)
        β > zero(β) || error("Laplace's scale must be positive")
        @compat new(Float64(μ), Float64(β))
    end

    @compat Laplace(μ::Real) = new(Float64(μ), 1.0)
    Laplace() = new(0.0, 1.0)
end

typealias Biexponential Laplace

@distr_support Laplace -Inf Inf

#### Parameters

location(d::Laplace) = d.μ
scale(d::Laplace) = d.β
params(d::Laplace) = (d.μ, d.β)

#### Statistics

mean(d::Laplace) = d.μ
median(d::Laplace) = d.μ
mode(d::Laplace) = d.μ

var(d::Laplace) = 2.0 * d.β^2
std(d::Laplace) = sqrt2 * d.β
skewness(d::Laplace) = 0.0
kurtosis(d::Laplace) = 3.0

entropy(d::Laplace) = log(2.0 * d.β) + 1.0

#### Evaluations

zval(d::Laplace, x::Float64) = (x - d.μ) / d.β
xval(d::Laplace, z::Float64) = d.μ + z * d.β

pdf(d::Laplace, x::Float64) = 0.5 * exp(-abs(zval(d, x))) / scale(d)
logpdf(d::Laplace, x::Float64) = - (abs(zval(d, x)) + log(2.0 * scale(d)))

cdf(d::Laplace, x::Float64) = (z = zval(d, x); z < 0.0 ? 0.5 * exp(z) : 1.0 - 0.5 * exp(-z))
ccdf(d::Laplace, x::Float64) = (z = zval(d, x); z > 0.0 ? 0.5 * exp(-z) : 1.0 - 0.5 * exp(z))
logcdf(d::Laplace, x::Float64) = (z = zval(d, x); z < 0.0 ? loghalf + z : loghalf + log2mexp(-z))
logccdf(d::Laplace, x::Float64) = (z = zval(d, x); z > 0.0 ? loghalf - z : loghalf + log2mexp(z))

quantile(d::Laplace, p::Float64) = p < 0.5 ? xval(d, log(2.0 * p)) : xval(d, -log(2.0 * (1.0 - p)))
cquantile(d::Laplace, p::Float64) = p > 0.5 ? xval(d, log(2.0 * (1.0 - p))) : xval(d, -log(2.0 * p))
invlogcdf(d::Laplace, lp::Float64) = lp < loghalf ? xval(d, logtwo + lp) : xval(d, -(logtwo + log1mexp(lp)))
invlogccdf(d::Laplace, lp::Float64) = lp > loghalf ? xval(d, logtwo + log1mexp(lp)) : xval(d, -(logtwo + lp))

function gradlogpdf(d::Laplace, x::Float64)
    μ, β = params(d)
    x == μ && error("Gradient is undefined at the location point")
    g = 1.0 / β
    x > μ ? -g : g
end

function mgf(d::Laplace, t::Real)
    st = d.β * t
    exp(t * d.μ) / ((1.0 - st) * (1.0 + st))
end
function cf(d::Laplace, t::Real)
    st = d.β * t
    cis(t * d.μ) / (1+st*st)
end

#### Sampling

function rand(d::Laplace)
    z = randexp()
    rand(Bool) ? d.μ + d.β * z : d.μ - d.β * z
end

#### Fitting

function fit_mle(::Type{Laplace}, x::Array)
    a = median(x)
    Laplace(a, mad(x, a))
end


immutable HermitianTensorOperator{S<:HilbertSpace,P,T}<:AbstractTensorMap{S,P,T}
    map::AbstractTensorMap{S,P,T}
    function HermitianTensorOperator(A::AbstractTensorMap{S,P,T})
        domain(A)==dual(conj(codomain(A))) || throw(SpaceError("Not an operator, i.e. domain != codomain"))
        new(A)
    end
end
hermitian{S,P,T}(A::AbstractTensorMap{S,P,T})=HermitianTensorOperator{S,P,T}(A::AbstractTensorMap{S,P,T})

# properties
domain(A::HermitianTensorOperator)=domain(A.map)
codomain(A::HermitianTensorOperator)=codomain(A.map)

Base.isreal(A::HermitianTensorOperator)=isreal(A.map)
Base.issym(A::HermitianTensorOperator)=isreal(A.map)
Base.ishermitian(A::HermitianTensorOperator)=true
Base.isposdef(A::HermitianTensorOperator)=isposdef(A.map)

# comparison
==(A::HermitianTensorOperator,B::HermitianTensorOperator)=A.map==B.map

# multiplication with vector
Base.A_mul_B!{S,P}(y::AbstractTensor{S,P},A::HermitianTensorOperator{S,P},x::AbstractTensor{S,P})=Base.A_mul_B!(y,A.map,x)
Base.Ac_mul_B!{S,P}(y::AbstractTensor{S,P},A::HermitianTensorOperator{S,P},x::AbstractTensor{S,P})=Base.A_mul_B!(y,A.map,x)


module TestDoctrines

using Test
using Catlab.Doctrines
using Catlab.Syntax

sexpr(expr::GATExpr) = sprint(show_sexpr, expr)
unicode(expr::GATExpr) = sprint(show_unicode, expr)
latex(expr::GATExpr) = sprint(show_latex, expr)

@testset "Categories" begin
  include("Category.jl")
end

@testset "MonoidalCategories" begin
  include("Monoidal.jl")
end

@testset "Permutations" begin
  include("Permutations.jl")
end

end


# Script to run through the files and fit the zero-inflated negative binomial distrubution

using Plots, Distributions, DelimitedFiles, Base.Printf
include("Utils.jl")
import Main.Utils
gr()

# Set some parameters
spread = 0.95
folder = "Chains/"
Files = readdir(folder)

# Make a stats variable to be writtten to a file
Stats = Array{Any,2}(undef,length(Files)+1,13)
Stats[1,:] = ["Name" "p" "w" "r" "K" "r_min" "r_max" "p_min" "p_max" "w_min" "w_max" "K_min" "K_max"]

for (ii,file) in enumerate(Files)

    # Load and plot the distribution data
    chain = readdlm(folder*file)

	# Extract parameters and plot
    r = Utils.find_MAP(chain,idx=1)
    p = Utils.find_MAP(chain,idx=2)
    w = Utils.find_MAP(chain,idx=3)
    K = Utils.find_MAP(chain,idx=4)
    int_r = Utils.credibleintervals(chain,idx=1, spread=spread)
    int_p = Utils.credibleintervals(chain,idx=2, spread=spread)
    int_w = Utils.credibleintervals(chain,idx=3, spread=spread)
    int_K = Utils.credibleintervals(chain,idx=4, spread=spread)

    # Plot and save
    plt = Utils.plot_chain(chain, [r,p,w,K], [int_r,int_p,int_w,int_K])
    try
        Plots.pdf("MCMC/"*file)
    catch
        mkdir("MCMC")
        Plots.pdf("MCMC/"*file)
    end

    # Save data
    println(file)
    Stats[ii+1,:] = [file p w r K r-int_r[1] int_r[2]-r p-int_p[1] int_p[2]-p w-int_w[1] int_w[2]-w K-int_K[1] int_K[2]-K]

end

DelimitedFiles.writedlm("Statistics.txt", Stats)


function factorial(x)
  if x ≤ 1
    return 1
  else
    return x * factorial(x - 1)
  end
end

v = factorial(big(9000))
s = string(v)

for c in s
  println(c)
end


using OrdinaryDiffEq, DiffEqSensitivity, LinearAlgebra, Optimization, OptimizationFlux, Flux
nn = Chain(Dense(1,16),Dense(16,16,tanh),Dense(16,2))
initial,re = Flux.destructure(nn)

function ode2!(u, p, t)
    f1, f2 = re(p)([t])
    [-f1^2; f2]
end

tspan = (0.0, 10.0)
prob = ODEProblem(ode2!, Complex{Float64}[0;0], tspan, initial)

function loss(p)
  sol = last(solve(prob, Tsit5(), p=p, sensealg=BacksolveAdjoint(autojacvec=ZygoteVJP())))
  return norm(sol)
end

optf = Optimization.OptimizationFunction((x,p) -> loss(x), Optimization.AutoZygote())
optprob = Optimization.OptimizationProblem(optf, initial)
res = Optimization.solve(optprob, ADAM(0.1), maxiters = 100)


using QuantumComputing

# Quantum circuit of Deutsch-Jozsa algorithm
#
#      +---+  +---+  +---+  +---+
# |0⟩ -+ X +--+ H +--+   +--+ H +--
#      +---+  +---+  |   |  +---+
#                    | U |
#             +---+  |   |
# |0⟩ --------+ H +--+   +---------
#             +---+  +---+
#

psi = Ket(0,0,0,0,0)
n = length(psi)

Uf = reduce(*,[Hadamard(n,i) for i=1:n])

G1 = Not(n,5)
G2 = reduce(*,[Hadamard(n,i) for i=1:n])
G3 = reduce(*,[Hadamard(n,i) for i=1:n-1])

psi = G3*Uf*G2*G1*psi

println(psi)
println(extract(psi,5))


inputs = [1.2, 5.1, 2.1]
weights = [3.1, 2.1, 8.7]
bias = 3.0

output = inputs[1]*weights[1] + inputs[2]*weights[2] + inputs[3]*weights[3] + bias
println(output)


using NonlinearEigenproblemsTest
using NonlinearEigenproblems
using Test
using LinearAlgebra
using SpecialFunctions
using SparseArrays

@testset "compute_types" begin
    A0=(1:5)*(1:5)'
    A1=(1:5)*(3:7)'+I;
    nep=DEP([A0,A1]);
    x=ones(5);
    s=compute_rf(ComplexF64,nep,x,
                 NonlinearEigenproblems.NEPSolver.ScalarNewtonInnerSolver(),λ=-0.5+3im)[1];

    @test abs(x'*compute_Mlincomb(nep,s,x)) < eps()*100;

    y=zeros(5); y[1]=3;
    s=compute_rf(ComplexF64,nep,x,
                 NonlinearEigenproblems.NEPSolver.ScalarNewtonInnerSolver(),λ=-0.5+3im,
                 y=y)[1];

    @test abs(y'*compute_Mlincomb(nep,s,x)) < eps()*100;
end


# Copyright 2014 Mitchell Kember. Subject to the MIT License.
# Google Code Jam Round 1A 2008
# Problem C: Numbers

# Returns the last three digits before the decimal point of the sum 3 + sqrt(5)
# raised to the nth power, as a 3-character string left-padded with zeros.
function last3(n)
	mat = [0 1; -4 6]^(n-1)
	x = (6mat[1,1] + 28mat[1,2] + 1_000_000) % 1000
	x = (x == 0) ? 999 : x - 1
	lpad(x, 3, 0)
end

function main()
	n = int(readline())
	for i = 1:n
		x = int(readline())
		println("Case #$i: ", last3(x))
	end
end

main()


using Tmatrix
using Zygote
using Plots
function objective_function(r_array_input, θ_array_input)
    # println(r_array)
    # println("I am <<objective_function>>, and the type of r_array is $(typeof(r_array))")
    # println("r_array = $r_array")

    # quadrupling_mesh
    r_θ_array = Tmatrix.quadruple_mesh_density(r_array_input, θ_array_input)
    r_array = r_θ_array[:, 1]
    θ_array = r_θ_array[:, 2]
    println(r_array)
    println(θ_array)
    ϕ_array = zeros(size(θ_array))

    T = Tmatrix.T_matrix_SeparateRealImag_arbitrary_mesh(
        n_max,
        wl_or_freq_input,
        input_unit,
        Eps_r_r_1,
        Eps_r_i_1,
        Mu_r_r_1,
        Mu_r_i_1,
        Eps_r_r_2,
        Eps_r_i_2,
        Mu_r_r_2,
        Mu_r_i_2,
        (collect(r_array)),
        θ_array,
        ϕ_array,
        rotationally_symmetric,
        symmetric_about_plane_perpendicular_z,
        BigFloat_precision,
    )
    k1_complex = Tmatrix.get_WaveVector(
        wl_or_freq_input;
        input_unit = input_unit,
        Eps_r = Complex(Eps_r_r_1, Eps_r_i_1),
        Mu_r = Complex(Mu_r_r_1, Mu_r_i_1),
    )
    # return abs(T[1]) # if we need to minimize the first element of T-matrix

    # return sum(abs.(T)) # if we need to minimize the sum of all real and imag

    return Tmatrix.get_OrentationAv_scattering_CrossSection_from_Tmatrix(T, k1_complex) # if we need to minimize scattering cross section

    """
    return -1 * Tmatrix.get_OrentationAv_emissivity_from_Tmatrix(
        T,
        Complex(k1_r, k1_i),
        Tmatrix.calculate_surface_area_of_axisymmetric_particle(r_array, θ_array)
    ) # if we need to minimize emissivity
    """
end

function ∂objective_function(r_array, θ_array)
    # println("I am <<∂objective_function>>, and the type of r_array is $(typeof(r_array))")
    # println("r_array = $r_array")
    return Zygote.gradient(
        r_array -> objective_function(r_array, θ_array),
        (collect(r_array)),
    )[1]
end

wl_or_freq_input = 1e-6;
input_unit = "m";
n_max = 1
Eps_r_r_1 = 1.0;
Eps_r_i_1 = 0.0;
Eps_r_r_2 = 1.5;
Eps_r_i_2 = 0.01;
Mu_r_r_1 = 1.0;
Mu_r_i_1 = 0.0;
Mu_r_r_2 = 1.0;
Mu_r_i_2 = 0.0;

n_θ_points = 10
θ_array = collect(LinRange(1e-6, pi - 1e-6, n_θ_points))
r_array = 0.5 * wl_or_freq_input * ones(size(θ_array))
rotationally_symmetric = true
symmetric_about_plane_perpendicular_z = false
BigFloat_precision = nothing

learning_rate = 0.5e-4
global loss_here = -1e6
global ∂loss = 0
println("Starting optimization for T-matrix ...")
loss_array = []
loss_old = 9999999999999
for n_iteration in 1:50
    global r_array = r_array .- learning_rate .* ∂loss

    global loss_here = objective_function(r_array, θ_array)
    global ∂loss = ∂objective_function(r_array, θ_array)
    append!(loss_array, loss_here)

    println()
    println("iteration #$n_iteration: loss_here = $loss_here, ∂loss = $∂loss")
    println("r_array = $r_array")

    xyz = vcat(
        Tmatrix.convert_coordinates_Sph2Cart.(r_array, θ_array, zeros(size(r_array)))...,
    )
    p1 = plot!(
        xyz[:, 1],
        xyz[:, 3],
        aspect_ratio = :equal,
        label = "particle after $n_iteration iterations",
        title = "scattering cross section = $loss_here m",
    )
    p2 = plot(
        1:length(loss_array),
        loss_array,
        xlabel = "iteration #",
        ylabel = "scattering cross section (m)",
    )
    fig = plot(p1, p2, layout = (1, 2), size = (1200, 800))
    mkpath("cache/iteration_particle_plots/maximizing_emissivity")
    savefig(
        fig,
        "cache/iteration_particle_plots/maximizing_emissivity/particle_geom_iteration_$(n_iteration).png",
    )
    # if (loss_here - loss_old) < 1e-6; break; end
    global loss_old = loss_here
end


function GasOutflow(Depth,Thp,GasPVT, Qmax; di=2.99, Ts=80, ∇T=1, ϵ=0.0006, Qn=10, Dn=50, tol=0.05, sg=0.6, Inc=0)

    # Create the range of Depth and Rates
    DepthRange=range(0, stop=Depth, length=Dn)
    if size(Qmax,1)==1
        Qrange=range(100, stop=Qmax, length=Qn)
    else
        Qrange=Qmax
        Qn=size(Qmax,1)
    end

    #Size of sensibilities
    Thpn=size(Thp,1)
    din=size(di,1)

        ## Gas
    Z=LinearInterpolation(GasPVT.P.*1, GasPVT.Z.*1)
    ρg=LinearInterpolation(GasPVT.P.*1, GasPVT.Rho.*1)
    βg=LinearInterpolation(GasPVT.P.*1, GasPVT.Bg.*1)
    μg=LinearInterpolation(GasPVT.P.*1, GasPVT.Mug.*1)
    Cg=LinearInterpolation(GasPVT.P.*1, GasPVT.Cg.*1)

    Ma=GasPVT.Rho[1]*10.73*(GasPVT.T[1]+460)/GasPVT.P[1]
    sg=Ma/28.96

        #Temperature
    T=LinearInterpolation([0, Depth], [Ts+460, 460+Ts+Depth*(∇T/100)])

    #create variables
    Pwf=zeros(Thpn*din,Qn)
    Pprofile=zeros(Dn,Qn,Thpn*din)
    angle=90 .- fill(Inc,Dn)
    leg=Array{String,1}(undef,Thpn*din)
       count=0
    for k=1:Thpn  #Loop for Thp
        for n=1:din #Loop for Diameter

            count=count+1
             P=zeros(Dn,Qn)      #Matrix of Pressures
            ∇P=zeros(Dn,Qn)      #Matriz of perssure gradient
            ∇Pnew=0.0
            P[1,:].=Thp[k]       #Initial pressure is Thp
            ∇P[1,:].=0.012       #Initial PRessure gradient

                #General Loop for each Rate
                for j=1:Qn

                    #General Loop for each depth
                    for i=2:Dn
                        e=tol+0.01
                        ∇Pguess=∇P[i-1,j]

                        while e>=tol

                            Pguess=∇Pguess*(DepthRange[i]-DepthRange[i-1])+P[i-1,j]
                            Nre=(4*28.97*sg*Qrange[j]*14.7)/(π*di[n]*μg(Pguess)*10.73*520)
                            f=(1/(-4*log((ϵ/3.7065)-(5.0452/Nre)*log((ϵ^1.1098/2.8257)+(7.149/Nre)^0.8981))))^2   #Friction factor
                            s=(-0.0375*sg*sin(angle[i]*π/180)*(DepthRange[i]-DepthRange[i-1]))/(Z(Pguess)*T[DepthRange[i]])
                            Pnew=sqrt(exp(-s)*P[i-1,j]^2-2.685e-3*((f*(Z(Pguess)*T[DepthRange[i]]*Qrange[j])^2 * (1-exp(-s)))/(sin(angle[i]*π/180)*di[n]^5)))
                            ∇Pnew=(Pnew-P[i-1,j])/(DepthRange[i]-DepthRange[i-1])
                            e=  abs(∇Pguess-∇Pnew)/∇Pnew
                            ∇Pguess=∇Pnew

                        end #end While
                    ∇P[i,j]=∇Pnew
                    P[i,j]=∇P[i,j]*(DepthRange[i]-DepthRange[i-1])+P[i-1,j]
                    end #End Loop Depth

                end #End Loop Rate
             Pwf[count,:].=P[end,:]
             Pprofile[:,:,count]=P
             leg[count]="Thp=$(Thp[k]); di=$(di[n])"
         end #End Loop Diameter

    end #End Loop Thp
    return DepthRange, Qrange, Pwf, leg, Pprofile
      end #End Function


using Revise, JPEGs, FileIO, Plots
using Test

img = load("mandrill.tiff")
i2 = JPEGs.rgb_to_ycbcr(img)
y,cb,cr = JPEGs.split_channels(i2)
downsampled_y = JPEGs.downsample_4_2_0(y)

test_subimg = [
[52 55 61 66 70 61 64 73];
[63 59 55 90 109 85 69 72];
[62 59 68 113 144 104 66 73];
[63 58 71 122 154 106 70 69];
[67 61 68 104 126 88 68 70];
[79 65 60 70 77 68 58 75];
[85 71 64 59 55 61 64 83];
[87 79 69 68 65 76 78 94];
]

test_block = JPEGs.Block(test_subimg)
dctd = JPEGs.dct(test_block)
quantized = JPEGs.quantize(dctd)
zigzagged = JPEGs.zigzag(quantized)

# test color conversion
@assert JPEGs.ycbcr_to_rgb(JPEGs.rgb_to_ycbcr(img)) == img

# test splitting and joining channels
@assert JPEGs.join_channels(JPEGs.split_channels(i2)) == i2

# test upsampling and downsampling
@assert JPEGs.downsample_4_2_0(JPEGs.upsample_4_2_0(downsampled_y)) == downsampled_y

# test splitting into blocks and rejoining
@assert JPEGs.join_blocks(JPEGs.split_into_blocks(y)) == y

# test dct and idct
@assert JPEGs.invdct(JPEGs.dct(test_block)).data == test_block.data

# Test quantizing and unquantizing
@assert JPEGs.quantize(JPEGs.unquantize(quantized)) == quantized

# Test zigzagging
@assert JPEGs.unzigzag(JPEGs.zigzag(quantized)) == quantized

# Test dc value to int and back
for i=-2047:2047
    @assert JPEGs.dc_code_to_value(JPEGs.value_to_dc_code(i)...) == i
end

# Test huffman string encoding of dc values
for i=0:11
    @assert JPEGs.huffman_string_to_length(JPEGs.length_to_huffman_string(i)) == i
end

# Test converting value to dc coded bitstring and back
for i=-2047:2047
    @assert JPEGs.bits_to_dc_value(JPEGs.value_to_dc_bits(i)) == i
end

# Test joining and splitting segments
soi = [0xFF, 0xD8]
something = [0xFF, 0xD4]
joined = [soi, something]
@assert JPEGs.split_into_segments(JPEGs.segments_into_bytes(joined)) == joined

bytearray = [0xAA, 0xBB, 0xCC, 0xDD]
@assert JPEGs.pack_bits(JPEGs.unpack_bits(bytearray)) == bytearray

testscan = BitArray([1,1,1,1,0,0,0,0])
@assert testscan == JPEGs.deserialize_into_jpeg_scan(JPEGs.serialize_jpeg_scan(testscan))

bytes = JPEGs.encode_image(img)
i3 = JPEGs.decode_image(bytes)

function timing()
    allbits = JPEGs.encode_image(img)
    i3 = JPEGs.decode_image(allbits)
end
#@benchmark timing()

plot(i3)


@testset "SBitstream" begin
    x = SBitstream(0.5)
    y = SBitstream(0.1)
    w = SBitstream(-0.1)
    X = SBitstream.(0.5 * rand(2, 2))
    Y = SBitstream.(-0.5 * rand(2, 2))
    v = SBitstream.([0.5 * rand(), -0.5 * rand()])

    @testset "FP Values" begin
        @test x.value == 0.5 && y.value == 0.1
        @test float(x) == x.value
        @testset for op in (+, -, *)
            @test float(op(x, y)) == op(float(x), float(y))
            @test float(op(x, w)) == op(float(x), float(w))
        end
        @test float(y / x) == float(y) / float(x)
        @test_logs (:warn, "SBitstream can only be ∈ [-1, 1] (saturation occurring).") float(x / y) == 1
        @test float(x ÷ 2) == float(x) / 2
        @test float(w ÷ 2) == float(w) / 2
        @test_throws ErrorException x ÷ 0.2
        @test float(sqrt(x)) == sqrt(float(x))
        @test float(decorrelate(x)) == float(x)
        @test float.(X * Y) == float.(X) * float.(Y)
        @test float.(X * v) == float.(X) * float.(v)
        @test float(norm(v)) == norm(float.(v))
        @test float(max(x, y)) == max(float(x), float(y))
    end

    @testset "Sample Generation" begin
        xbits = generate(x, 100000)
        @test mean(map(xbit -> pos(xbit) - neg(xbit), xbits)) ≈ x.value rtol = 0.1
    end

    @testset "Bit-Level Simulation" begin
        T = 50000
        @testset for op in (+, -, *, max)
            sim = simulatable(op, x, y)
            @test begin
                z = 0
                for t in 1:T
                    bit = pop!(sim(op, x, y))
                    z += pos(bit) - neg(bit)
                end

                z / T
            end ≈ float(op(x, y)) rtol = 0.1
            sim = simulatable(op, y, x)
            @test begin
                z = 0
                for t in 1:T
                    bit = pop!(sim(op, y, x))
                    z += pos(bit) - neg(bit)
                end

                z / T
            end ≈ float(op(y, x)) rtol = 0.1
            sim = simulatable(op, x, w)
            @test begin
                z = 0
                for t in 1:T
                    bit = pop!(sim(op, x, w))
                    z += pos(bit) - neg(bit)
                end

                z / T
            end ≈ float(op(x, w)) rtol = 0.1
        end
        @testset "op = ÷" begin
            sim = simulatable(÷, x, 2)
            @test begin
                z = 0
                for t in 1:T
                    bit = pop!(sim(÷, x, 2))
                    z += pos(bit) - neg(bit)
                end

                z / T
            end ≈ float(x ÷ 2) rtol = 0.1
            sim = simulatable(÷, w, 2)
            @test begin
                z = 0
                for t in 1:T
                    bit = pop!(sim(÷, w, 2))
                    z += pos(bit) - neg(bit)
                end

                z / T
            end ≈ float(w ÷ 2) rtol = 0.1
        end
        @testset "op = /" begin
            sim = simulatable(/, y, x)
            @test begin
                z = 0
                for t in 1:T
                    bit = pop!(sim(/, y, x))
                    z += pos(bit) - neg(bit)
                end

                z / T
            end ≈ float(y / x) rtol = 0.1
        end
        @testset for op in (sqrt, decorrelate)
            sim = simulatable(op, x)
            @test begin
                z = 0
                for t in 1:T
                    bit = pop!(sim(op, x))
                    z += pos(bit) - neg(bit)
                end

                z / T
            end ≈ float(op(x)) rtol = 0.1
        end
        @testset "op = max(x, y, z)" begin
            sim = simulatable(max, x, y, w)
            @test begin
                z = 0
                for t in 1:T
                    bit = pop!(sim(max, x, y, w))
                    z += pos(bit) - neg(bit)
                end

                z / T
            end ≈ float(max(x, y, w)) rtol = 0.1
        end
        @testset "op = * (matrix, matrix)" begin
            sim = simulatable(*, X, Y)
            @test begin
                Z = zeros(2, 2)
                for t in 1:T
                    bit = pop!.(sim(*, X, Y))
                    Z .+= pos.(bit) .- neg.(bit)
                end

                Z ./ T
            end ≈ float.(X * Y) rtol = 0.1
            sim = simulatable(*, Y, X)
            @test begin
                Z = zeros(2, 2)
                for t in 1:T
                    bit = pop!.(sim(*, Y, X))
                    Z .+= pos.(bit) .- neg.(bit)
                end

                Z ./ T
            end ≈ float.(Y * X) rtol = 0.1
        end
        @testset "op = * (matrix, vector)" begin
            sim = simulatable(*, X, v)
            @test begin
                T = 100000
                Z = zeros(2)
                for t in 1:T
                    bit = pop!.(sim(*, X, v))
                    Z .+= pos.(bit) .- neg.(bit)
                end

                Z ./ T
            end ≈ float.(X * v) rtol = 0.1
            sim = simulatable(*, transpose(v), X)
            @test begin
                T = 100000
                Z = zeros(1, 2)
                for t in 1:T
                    bit = pop!.(sim(*, transpose(v), X))
                    Z .+= pos.(bit) .- neg.(bit)
                end

                Z ./ T
            end ≈ float.(transpose(v) * X) rtol = 0.1
        end
        @testset "op = norm" begin
            sim = simulatable(norm, v)
            @test begin
                T = 100000
                z = 0
                for t in 1:T
                    bit = pop!(sim(norm, v))
                    z += pos(bit) - neg(bit)
                end

                z / T
            end ≈ float(norm(v)) rtol = 0.1
        end
    end
end


using Distances
using NearestNeighbors
import StateSpaceReconstruction: cembed

@testset "Kraskov (kNN) 3D $i" for i in 1:n_realizations
    # Some random points that we embed
    pts = rand(3, 100)
    E = cembed(pts)

    # How many neighbors to consider?
    k1 = 2
    k2 = 3
    # Which variables go into which marginals?
    target_future = [1]
    target_presentpast = [2]
    source_presentpast = [3]
    conditioned_presentpast = Int[]

    v = TEVars(target_future,
            target_presentpast,
            source_presentpast,
            conditioned_presentpast)

    @testset "On raw points, default metric" begin
        @test transferentropy_kraskov(pts, k1, k2, v) |> isfinite
        @test transferentropy_kraskov(pts, k1, k2, target_future,
                                                    target_presentpast,
                                                    source_presentpast,
                                                conditioned_presentpast) |> isfinite
    end

    @testset "On embeddings, default metric" begin
        @test transferentropy_kraskov(E, k1, k2, v) |> isfinite
        @test transferentropy_kraskov(E, k1, k2, target_future,
                                                    target_presentpast,
                                                    source_presentpast,
                                                    conditioned_presentpast) |> isfinite
    end

    @testset "On points, specifying the distance metric" begin
        @test transferentropy_kraskov(pts, k1, k2, v; metric = Chebyshev()) |> isfinite
        @test transferentropy_kraskov(pts, k1, k2, v; metric = Chebyshev()) |> isfinite
    end

    @testset "On embeddings, specifying the distance metric" begin
        @test transferentropy_kraskov(pts, k1, k2, v; metric = Chebyshev()) |> isfinite
        @test transferentropy_kraskov(pts, k1, k2, v; metric = Chebyshev()) |> isfinite
        @test transferentropy_kraskov(E, k1, k2, v; metric = Chebyshev()) |> isfinite
        @test transferentropy_kraskov(E, k1, k2, v; metric = Chebyshev()) |> isfinite
    end
end

@testset "Kraskov (kNN) 4D $i" for i in 1:n_realizations
    # Some random points that we embed
    pts = rand(4, 200)
    E = cembed(pts)

    # How many neighbors to consider?
    k1 = 2
    k2 = 3
    # Which variables go into which marginals?
    target_future = [1]
    target_presentpast = [2, 3]
    source_presentpast = [4]
    conditioned_presentpast = Int[]

    v = TEVars(target_future,
            target_presentpast,
            source_presentpast,
            conditioned_presentpast)

    @testset "On raw points, default metric" begin
        @test transferentropy_kraskov(pts, k1, k2, v) |> isfinite
        @test transferentropy_kraskov(pts, k1, k2, target_future,
                                                    target_presentpast,
                                                    source_presentpast,
                                                conditioned_presentpast) |> isfinite
    end

    @testset "On embeddings, default metric" begin
        @test transferentropy_kraskov(E, k1, k2, v) |> isfinite
        @test transferentropy_kraskov(E, k1, k2, target_future,
                                                    target_presentpast,
                                                    source_presentpast,
                                                    conditioned_presentpast) |> isfinite
    end

    @testset "On points, specifying the distance metric" begin
        @test transferentropy_kraskov(pts, k1, k2, v; metric = Chebyshev()) |> isfinite
        @test transferentropy_kraskov(pts, k1, k2, v; metric = Chebyshev()) |> isfinite
    end

    @testset "On embeddings, specifying the distance metric" begin
        @test transferentropy_kraskov(pts, k1, k2, v; metric = Chebyshev()) |> isfinite
        @test transferentropy_kraskov(pts, k1, k2, v; metric = Chebyshev()) |> isfinite
        @test transferentropy_kraskov(E, k1, k2, v; metric = Chebyshev()) |> isfinite
        @test transferentropy_kraskov(E, k1, k2, v; metric = Chebyshev()) |> isfinite
    end
end