# rootsah for BigInt and BigFloat
function rootsah(pol::Union{Poly{BigInt},Poly{BigFloat}}, D::Vector{Float64})
    # COMPUTES: the roots of polynomials with all distinct real roots.
    # The computation is forward stable. The program uses SymArrow (arrowhead) companion matrix and
    # corresponding eig routine
    # D are barycentric coordinates - elements od D must interpolate the roots of P,
    # for example
    #              D=roots(polyder(pol))
    # RETURNS: roots E

    # Type is BigFloat
    T = BigFloat

    # τ = [1e2,1e2,1e2,1e2,1e2] or similar is the vector of tolerances for eig
    τ=[1e2,1e2,1e2,1e2,1e2]

    Dm=map(T,D)
    Dm=sort(Dm,rev=true)
    D=sort(D,rev=true)
    p=map(T,[pol[i] for i=0:1:length(pol)-1])
    p=p[end:-1:1]

    # Compute z of the arrowhead
    # First we compute values s=p(D) we use Horner sheme with double the working precision
    setprecision(512)
    n=length(p)-1
    s=Array{BigFloat}(undef,n-1)
    for i=1:n-1
        s[i]=one(T)
    end
    s=p[1]*s
    for i=2:n+1
        r=s.*Dm
        s=r.+p[i]
    end

    # Compute t's
    t=Array{BigFloat}(undef,n-1)
    for j=1:n-1
        h=one(T)
        for i=1:j-1
            g=Dm[j]-Dm[i]
            h=h*g
        end
        for i=j+1:n-1
            g=Dm[j]-Dm[i]
            h=h*g
        end
        t[j]=h
    end
    # Compute alphaD
    αD=p[2]/p[1]
    for i=1:n-1
        αD+=Dm[i]
    end
     αD=-αD
    #  Compute z
    zD=Array{BigFloat}(undef,n-1)
    for i=1:n-1
        zD[i]=sqrt((-s[i])/(t[i]*p[1]))
    end

    z=Array{Float64}(undef,n-1)
    for i=1:n-1
        z[i]=Float64(zD[i])
    end
    α=Float64(αD)
    A=SymArrow(D,z,α,n)
    E=zeros(Float64,n)
    Qout=zeros(Int64,n)
    for k=1:n
        E[k], Qout[k] = eigen( A,zD,αD,k,τ )
    end

    # reset the bigfloat_precision
    setprecision(256)
    E, Qout
end

#------------------------
#--------  Functions

#-------- Inverses

function inv(A::SymArrow{T},zD::Vector{BigFloat}, αD::BigFloat,i::Int64,τ::Vector{Float64}=[1e3, 10*length(A.D)]) where T
    # COMPUTES: inverse of a SymArrow matrix A, inv(A-A.D[i]*I) which is again SymArrow
    # uses higher precision to compute top of the arrow element accurately, if
    # needed.
    # τ=[tolb,tolz] are tolerances, usually [1e3, 10*n]
    # [0.0,0.0] forces DoubleDouble, [1e50,1e50] would never use it
    # RETURNS:  SymArrow(D,z,b,i), Kb, Kz, Qout
    # Kb - condition Kb, Kz - condition Kz, Qout = 1 / 0 - double was / was not used

    n=length(A.D)
    D=Array{T}(undef,n)
    z=Array{T}(undef,n)
    wz=1/A.z[i]
    σ=A.D[i]
    for k=1:i-1
        D[k]=one(T)/(A.D[k]-σ)
        z[k]=-A.z[k]*D[k]*wz
    end
    for k=i+1:n
        D[k-1]=one(T)/(A.D[k]-σ)
        z[k-1]=-A.z[k]*D[k-1]*wz
    end
    a=A.a-A.D[i]

    # compute the sum in a plain loop
    P=zero(T)
    Q=zero(T)
    Qout=0
    nn=n-1
    for k=1:i-1
        D[k]>0.0 ? P+=A.z[k]^2*D[k] : Q+=A.z[k]^2*D[k]
    end
    for k=i:nn
        D[k]>0.0 ? P+=A.z[k+1]^2*D[k] : Q+=A.z[k+1]^2*D[k]
    end
    a<0 ? P-=a : Q-=a

    Kb=(P-Q)/abs(P+Q)
    Kz=maximum(abs,A.z)*abs(wz)
    if Kb<τ[1] ||  Kz<τ[2]
        b=(P+Q)*wz*wz
    else  # recompute in double the working precision
        Qout=1
        AD=map(BigFloat,A.D)
        σd=AD[i]
        Dd=[[AD[k]-σd for k=1:i-1];
            [AD[k]-σd for k=i+1:length(A.D)]]
        wzd=zD[i] #
        ad=αD-σd
        Pd,Qd=map(BigFloat,(0.0,0.0))
        for k=1:i-1
            Dd[k] > 0.0 ? Pd+=zD[k]^2/Dd[k] :
            Qd+=zD[k]^2/Dd[k]
        end
        for k=i:nn
            Dd[k] > 0.0 ? Pd+=zD[k+1]^2/Dd[k] :
            Qd+=zD[k+1]^2/Dd[k]
        end
        ad<0 ?   Pd-=ad : Qd-=ad
        bd=(Pd+Qd)/(wzd*wzd)
        b=Float64(bd)
    end

    if i<A.i
        # SymArrow([D[1:A.i-2],zero(T),D[A.i-1:end]],[z[1:A.i-2],wz,z[A.i-1:end]],b,i),Kb,Kz,Qout
        # This reduces memory allocation (4 vectors) and is 30-40% faster
        for k=n-1:-1:A.i-1
            D[k+1]=D[k]
            z[k+1]=z[k]
        end
        D[A.i-1]=zero(T)
        z[A.i-1]=wz
        SymArrow(D,z,b,i),Kb,Kz,Qout
    else
        # SymArrow([D[1:A.i-1],zero(T),D[A.i:end]],[z[1:A.i-1],wz,z[A.i:end]],b,i+1),Kb,Kz,Qout
        for k=n-1:-1:A.i
            D[k+1]=D[k]
            z[k+1]=z[k]
        end
        D[A.i]=zero(T)
        z[A.i]=wz
        SymArrow(D,z,b,i+1),Kb,Kz,Qout
    end
end # inv

function inv(A::SymArrow{T}, zD::Vector{BigFloat}, αD::BigFloat, σ::Float64, τ::Float64=1.0e3) where T
    # COMPUTES: inverse of the shifted SymArrow A, inv(A-shift*I) which is SymDPR1
    # uses DoubleDouble to compute ρ accurately, if needed.
    # τ is tolerance, usually 1e3,  0.0 forces Double, 1e50 would never use it
    # RETURNS: SymDPR1(D,u,ρ), Kρ, Qout
    # Kρ - condition Kρ, Qout = 1 / 0 - Double was / was not used

    n=length(A.D)
    D=Array{T}(undef,n+1)
    u=Array{T}(undef,n+1)
    for k=1:A.i-1
        D[k]=one(T)/(A.D[k]-σ)
        u[k]=A.z[k]*D[k]
    end
    D[A.i]=zero(T)
    u[A.i]=-one(T)
    for k=A.i:n
        D[k+1]=one(T)/(A.D[k]-σ)
        u[k+1]=A.z[k]*D[k+1]
    end

    # compute ρ and Kρ
    #--- compute the sum in a plain loop
    a=A.a-σ
    P=zero(T)
    Q=zero(T)
    Qout=0
    for k=1:A.i-1
        D[k]>0.0 ? P+=A.z[k]^2*D[k] : Q+=A.z[k]^2*D[k]
    end
    for k=A.i:n
        D[k+1]>0.0 ? P+=A.z[k]^2*D[k+1] : Q+=A.z[k]^2*D[k+1]
    end
    a<0 ? P-=a : Q-=a

    # Condition of ρ
    Kρ=(P-Q)/abs(P+Q)
    if Kρ<τ
        ρ=-1.0/(P+Q)
    else  # recompute in double the working precision
        Qout=1
        Pd,Qd=map(BigFloat,(0.0,0.0))
        AD=map(BigFloat,A.D)
        σd=map(BigFloat,σ)
        ad=αD-σd
        for k=1:A.i-1
            D[k]>0.0 ? Pd+=zD[k]^2/(AD[k]-σd) : Qd+=zD[k]^2/(AD[k]-σd)
        end
        for k=A.i:n
            D[k+1]>0.0 ? Pd+=zD[k]^2/(AD[k]-σd) : Qd+=zD[k]^2/(A.D[k]-σd)
        end
        ad<0 ? Pd-=ad : Qd-=ad
        r=map(BigFloat,1.0)/(Pd+Qd)
        ρ=Float64(-r)
    end

    # returns the following
    SymDPR1(D,u,ρ), Kρ, Qout
end # inv

function inv(A::SymArrow{T}, zD::Vector{BigFloat}, αD::BigFloat, σ::BigFloat) where T
    # COMPUTES: inverse of the shifted SymArrow A, inv(A-σ*I), which is a SymDPR1
    # here σ is BigFloat so it uses double the working precision to compute everything
    # RETURNS: SymDPR1(D1,u1,ρ), Qout
    # Qout = 1 on exit meaning Double was used

    n=length(A.D)
    D=Array{BigFloat}(undef,n+1)
    u=Array{BigFloat}(undef,n+1)
    for k=1:n
        D[k]=map(BigFloat,A.D[k])-σ
    end
    u[1:n]=zD
    a=αD-σ
    i=A.i

    oned=one(BigFloat)
    zerod=zero(BigFloat)
    for k=1:i-1
        u[k]=u[k]/D[k]
        D[k]=oned/D[k]
    end

    for k=i:n
        u[k+1]=u[k]/D[k]
        D[k+1]=oned/D[k+1]
    end
    D[i]=zerod
    u[i]=-oned

    # compute ρ and Kρ
    #--- compute the sum in a plain loop
    P,Q=zerod,zerod
    Qout=1
    for k=1:i-1
        D[k] > 0.0 ? P+=zD[k]*u[k] : Q+=zD[k]*u[k]
    end
    for k=i:n
        D[k+1] > 0.0 ? P+=zD[k]*u[k+1] : Q+=zD[k]*u[k+1]
    end
    a<0.0  ? P-=a : Q-=a
    r=oned/(P+Q)
    ρ=Float64(-r)

    # returns the following
    D₁=Array{T}(undef,n+1)
    u₁=Array{T}(undef,n+1)
    for k=1:n+1
        D₁[k]=Float64(D[k])
    end
    for k=1:n+1
        u₁[k]=Float64(u[k])
    end
    SymDPR1(D1,u1,ρ), Qout
end # inv

function  eigen( A::SymArrow{T},zD::Vector{BigFloat},αD::BigFloat,k::Int64,
    τ::Vector{Float64}=[1e3,10.0*length(A.D),1e3,1e3,1e3]) where T
    # COMPUTES: k-th eigenvalue of an ordered irreducible SymArrow
    # A = [Diagonal(D) z; z' alpha]
    # Specially designed to be used in the polynomial rootfinder rootsah !!!!
    # τ=[tolb,tolz,tolnu,tolrho,tollambda] = [1e3,10.0*n,1e3,1e3,1e3]
    # RETURNS: λ, Qout
    # λ - k-th eigenvalue
    # Kb, Kz, Kν, Kρ - condition numbers
    # Qout = 1 / 0 - Double was / was not used

    # If Qout>0, quad precision was used
    # i was the shift index

    # Set the dimension
    n = length(A.D) + 1
    # Set all conditions initially to zero
    Kb,Kz,Kν,Kρ=0.0,0.0,0.0,0.0
    Qout=0
    v=zeros(T,n)
    # Kz is former kappa_nu

    # Determine the shift σ, the shift index i, and whether λ
    # is on the left or the right side of the nearest pole
    # Exterior eigenvalues (k = 1 or k = n):
    if k == 1
        σ,i,side = A.D[1],1,'R'
    elseif k==n
        σ,i,side = A.D[n-1],n-1,'L'
    else
        # Interior eigenvalues (k in (2,...n-1) ):
        # We need to compute Fmiddle carefully
        Dtemp = A.D.-A.D[k]
        middle = Dtemp[k-1]/2.0
        P=zero(Float64)
        Q=zero(Float64)
        P+=sum(A.z[1:k-1].^2 ./(Dtemp[1:k-1].-middle))
        Q+=sum(A.z[k:n-1].^2 ./(Dtemp[k:n-1].-middle))

        P₁=map(BigFloat,P)
        Q₁=map(BigFloat,Q)
        atemp = (αD-map(BigFloat,A.D[k]))-map(BigFloat,middle)
        atemp<0 ? P₁-=atemp : Q₁-=atemp
        # Fmiddle = (atemp-middle)-sum(zD.^2./(Dtemp-middle))
        Fmiddle=-(P₁+Q₁)
        σ,i,side = Fmiddle < zero(BigFloat) ? (A.D[k],k,'R') : (A.D[k-1],k-1,'L')
    end

    # Compute the inverse of the shifted matrix, A_i^(-1), Kb and Kz
    Ainv,Kb,Kz,Qout = inv(A,zD,αD,i)

    # Compute the eigenvalue of the inverse shifted matrix
    ν = bisect( Ainv,side )

    #  ν=fastzero([invD1; 0; invD2], [w1;wz;w2], b, side); # implement later
    #  [ν-ν₁]/ν, [ν-nueva]/ν, pause  # and compare
    if abs(ν)==Inf
        # this is nonstandard
        # Deflation in aheig (ν=Inf)
        v[i]=1.0
        λ=σ
    else
        # standard case, full computation
        # ν₁ is the F- or 1-norm of the inverse of the shifted matrix
        # ν₁0=maximum([sum(abs(Ainv.z))+abs(Ainv.a), maximum(abs(Ainv.D)+abs(Ainv.z))])
        ν₁=0.0
        for l=1:n-1
            ν₁=max(ν₁,abs(Ainv.D[l])+abs(Ainv.z[l]))
        end
        ν₁=max(sum(abs,Ainv.z)+abs(Ainv.a), ν₁)
        Kν=ν₁/abs(ν)
        if Kν<τ[3]
            # Accuracy is fine, compute the eigenvector
            μ = 1.0/ν
            # v=[ A.z./((A.D-σ)-μ);-1.0]
            for l=1:n-1
            	v[l] = A.z[l]/((A.D[l]-σ)-μ)
            end
            v[n]=-1.0
            λ = μ+σ
            normalize!(v)
        else
            # Remedies according to Remark 3 - we shift between original
            # eigenvalues and compute DPR1 matrix
            # 1/ν₁+σ, 1/ν+σ
            # println("Remedy 3 ")
            ν = side=='R' ? abs(ν) : -abs(ν)
            ν₁=-sign(ν)*ν₁
            σ₁=(ν₁+ν)/(2.0*ν*ν₁)+σ
            if count(isequal(σ₁),A.D)>0 # we came back with a pole
                # recompute σᵥ more accurately according with dekker
                setprecision(512)
                σᵥ=(ν₁+ν)/(BigFloat(2.0)*ν*ν₁)+σ
                # Compute the inverse of the shifted arrowhead (DPR1)
                Ainv,Qout1=inv(A,σᵥ) # Ainv is Float64
                ν₁=bisect(Ainv,side)
                μ₁ = 1.0/ν₁
                D0=map(BigFloat,A.D)-σᵥ
                D1=[D0[l] for l=1:length(D0)]
                if count(isequal(μ₁),D1)>0
                    ind=findall(isequal(μ₁),D1)
                    v=zeros(n)
                    v[ind]=1.0
                else
                    v=[ A.z./(D1-μ₁);-1.0]
                end
                # Shift the eigenvalue back in Double
                lam = BigFloat(1.0)/ν₁+σᵥ
                # Return this
                λ = lam
                normalize!(v)
            else
                # Compute the inverse of the shifted arrowhead (DPR1)
                Ainv, Kρ,Qout₁=inv(A,σ₁,τ[4]) # Ainv is Float64
                # Compute the eigenvalue by bisect for DPR1
	            # Note: instead of bisection could use dlaed4 (need a wrapper) but
	            # it is not faster. There norm(u)==1
                ν₁= bisect(Ainv,side)
                μ₁=1.0/ν₁
                # standard v
                v=[ A.z./((A.D-σ₁)-μ₁);-1.0]
                # Return this - shift the eigenvalue back and normalize the vector
                λ = μ₁+σ₁
                normalize!(v)
            end
            Qout==1 && (Qout=Qout+2)
        end

        # Remedy according to Remark 1 - we recompute the the eigenvalue
        # near zero from the inverse of the original matrix (a DPR1 matrix).
        if (abs(A.D[i])+abs(1.0/ν))/abs(λ)>τ[5]
            if k==1 && A.D[1]<0.0 || k==n && A.D[n-1]>0.0 || i<n-1 && side=='L' && sign(A.D[i])+sign(A.D[i+1])==0 ||
                side=='R' && sign(A.D[i])+sign(A.D[i-1])==0
                # println("Remedy 1 ")
                # Compute the inverse of the original arrowhead (DPR1)
                Ainv,Kρ,Qout₁ = inv(A,0.0,τ[4]) # Ainv is Float64
                Qout==1 && (Qout=Qout+4)
                if isinf(Ainv.r) || isnan(Ainv.r)
                    λ=0.0
                else
                    # Here we do not need bisection. We compute the Rayleigh
                    # quotient by using already computed vectors which is
                    # componentwise accurate
                    ν₁=sum(v.^2 .*Ainv.D)+Ainv.r*sum(v.*Ainv.u)^2;
                    λ=1.0/ν₁
                end
            end
        end
    end

    # Return this
    Float64(λ), Qout
end # eigen (k)

function rootsWDK(p::Poly{T},x₀::Vector{T1},steps) where {T,T1}
    # Polynomial roots with Weierstrass, Durand, Kerner method
    n=length(x₀)
    x=x₀
    for k=1:steps
        for i=1:n
            x[i]=x[i]-polyval(p,x[i])/(p[end]*prod(x[i].-x[collect([1:i-1;i+1:n])]))
        end
    end
    x
end


using Distributed
@everywhere using MCPhylo, LinearAlgebra

## Data
jaws = Dict{Symbol, Any}(
  :Y =>
   [47.8 48.8 49.0 49.7
    46.4 47.3 47.7 48.4
    46.3 46.8 47.8 48.5
    45.1 45.3 46.1 47.2
    47.6 48.5 48.9 49.3
    52.5 53.2 53.3 53.7
    51.2 53.0 54.3 54.5
    49.8 50.0 50.3 52.7
    48.1 50.8 52.3 54.4
    45.0 47.0 47.3 48.3
    51.2 51.4 51.6 51.9
    48.5 49.2 53.0 55.5
    52.1 52.8 53.7 55.0
    48.2 48.9 49.3 49.8
    49.6 50.4 51.2 51.8
    50.7 51.7 52.7 53.3
    47.2 47.7 48.4 49.5
    53.3 54.6 55.1 55.3
    46.2 47.5 48.1 48.4
    46.3 47.6 51.3 51.8],
  :age => [8.0, 8.5, 9.0, 9.5]
)
M = jaws[:M] = size(jaws[:Y], 2)
N = jaws[:N] = size(jaws[:Y], 1)
jaws[:y] = vec(jaws[:Y])
jaws[:x] = kron(ones(jaws[:N]), jaws[:age])

## Model Specification
model = Model(

  y = Stochastic(1,
    (beta0, beta1, x, Sigma) -> BDiagNormal(beta0 .+ beta1 * x, Sigma),
    false
  ),

  beta0 = Stochastic(
    () -> Normal(0, sqrt(1000))
  ),

  beta1 = Stochastic(
    () -> Normal(0, sqrt(1000))
  ),

  Sigma = Stochastic(2,
    M -> InverseWishart(4.0, Matrix{Float64}(I, M, M))
  )

)

## Initial Values
inits = [
  Dict(:y => jaws[:y], :beta0 => 40, :beta1 => 1, :Sigma => Matrix{Float64}(I, M, M)),
  Dict(:y => jaws[:y], :beta0 => 10, :beta1 => 10, :Sigma => Matrix{Float64}(I, M, M))
]

## Sampling Scheme
scheme = [Slice([:beta0, :beta1], [10, 1]),
          AMWG(:Sigma, 0.1)]
setsamplers!(model, scheme)

## MCMC Simulations
sim = mcmc(model, jaws, inits, 10000, burnin=2500, thin=2, chains=2)
describe(sim)


"""
Simplex Crossover (SPX).

`ϵ>0` controls how the original simplex is inflated, `ϵ=1` means no
inflation.

See Tsutsui, Yamamura & Higuchi "Multi-parent recombination with simplex crossover in real coded genetic algorithms", 1999,
Proc. of the Genetic and Evolutionary Computation Conference
"""
struct SimplexCrossover{NP} <: CrossoverOperator{NP,1}
    ϵ::Float64      # inflation rate

    SimplexCrossover{NP}(ϵ::Number = Math.sqrt(NP + 1)) where NP = new{NP}(ϵ)
    SimplexCrossover{NP}(params::Parameters) where NP = new{NP}(params[:SPX_ϵ])
end

const SPX_DefaultOptions = ParamsDict(
    :SPX_ϵ => 1.1 # how much to inflate
)

#masscenter(pop, parentIndices) = mapslices(mean, pop[:, parentIndices], 2)

function apply!(xover::SimplexCrossover,
                target::AbstractIndividual, targetIndex::Int,
                pop, parentIndices)
    @assert length(parentIndices) == numparents(xover)

    offset = fill!(similar(target), 0)
    parentSum = copy(viewer(pop, parentIndices[1]))
    for i in 1:(numparents(xover)-1)
        s = rand()^(1/i)
        parent = viewer(pop, parentIndices[i])
        parentNext = viewer(pop, parentIndices[i+1])
        @inbounds offset .= s * (xover.ϵ * (parent .- parentNext) .+ offset)
        @inbounds parentSum .+= parentNext
    end
    s = (1-xover.ϵ)/numparents(xover)
    parentLast = viewer(pop, parentIndices[end])
    @inbounds target .= s .* parentSum + xover.ϵ .* parentLast .+ offset
    return target
end


abstract type MultivariateFun{T,N} <: Function end
const BivariateFun{T} = MultivariateFun{T,2}

export grad, lap, curl

#implements coefficients/values/evaluate
space(f::MultivariateFun{T,N}) where {T,N}=mapreduce(k->space(f,k),⊗,1:N)
domain(f::MultivariateFun{T,N}) where {T,N}=mapreduce(k->domain(f,k),×,1:N)

domain(f::MultivariateFun,k::Integer)=domain(space(f,k))

differentiate(u::BivariateFun,i::Integer,j::Integer) =
    j==0 ? u : differentiate(differentiate(u,i),i,j-1)
grad(u::BivariateFun) = [differentiate(u,1),differentiate(u,2)]
lap(u::BivariateFun) = differentiate(u,1,2)+differentiate(u,2,2)
Base.div(u::AbstractVector{B}) where {B<:BivariateFun} =
    differentiate(u[1],1)+differentiate(u[2],2)
curl(u::AbstractVector{B}) where {B<:BivariateFun} = differentiate(u[2],1)-differentiate(u[1],2)

Base.chop(f::MultivariateFun) = chop(f,10eps())
cfstype(::MultivariateFun{T}) where {T} = T
cfstype(::Type{MultivariateFun{T,N}}) where {T,N} = T
cfstype(::Type{MF}) where {MF<:MultivariateFun} = cfstype(supertype(MF))

include("VectorFun.jl")
include("TensorSpace.jl")
include("LowRankFun.jl")
include("ProductFun.jl")

arglength(f)=length(Base.uncompressed_ast(f.code.def).args[1])

## Convert between Fun and MultivariateFun
# need to chop extra zeros
Fun(f::ProductFun) =
    Fun(space(f),chop!(fromtensor(space(f),coefficients(f)),0))
Fun(f::ProductFun,sp::TensorSpace) = Fun(ProductFun(f,sp))
Fun(f::LowRankFun) = Fun(ProductFun(f))

Fun(f::MultivariateFun,sp::Space) = Fun(Fun(f),sp)

Fun(f,d1::Domain,d2::Domain) = Fun(f,d1*d2)

coefficients(f::BivariateFun,sp::TensorSpace)=coefficients(f,sp[1],sp[2])

points(f::BivariateFun,k...)=points(space(f),size(f,1),size(f,2),k...)

function *(vx::LowRankFun,u0::ProductFun)
    ret=zero(space(u0))
    for k=1:length(vx.A)
        a,b=vx.A[k],vx.B[k]
        ret+=transpose(b*(transpose(a*u0)))
    end
    ret
end

*(a::ProductFun,b::LowRankFun)=b*a
*(a::MultivariateFun,b::MultivariateFun)=LowRankFun(a)*ProductFun(b)

for OP in (:+,:-,:*,:/)
    @eval begin
        $OP(f::Fun,g::MultivariateFun)=$OP(ProductFun(f),g)
        $OP(f::MultivariateFun,g::Fun)=$OP(f,ProductFun(g))
    end
end

Base.sum(f::Fun{TS},k::Integer) where {TS<:TensorSpace}=sum(ProductFun(f),k)
Base.sum(f::Fun{TS}) where {TS<:TensorSpace}=sum(ProductFun(f))

## kron
# TODO: generalize
function Base.kron(f::Fun,g::Fun)
    sp=space(f)⊗space(g)
    it=tensorizer(sp)
    N=ncoefficients(f);M=ncoefficients(g)
    cfs=Array{promote_type(cfstype(f),cfstype(g))}(undef,0)
    for (k,j) in it
        # Tensor product is N x M, so if we are outside
        # the (N+M)th diagonal we have no more entries
        if k+j > N+M
            break
        elseif k ≤ N && j ≤ M
            push!(cfs,f.coefficients[k]*g.coefficients[j])
        else
            push!(cfs,0)
        end
    end
    Fun(sp,cfs)
end
Base.kron(f::Fun,g::Number) = kron(f,Fun(g))
Base.kron(f::Number,g::Fun) = kron(Fun(f),g)


using DataFrames
using CSV
using Plots
using Unitful
using UnitfulRecipes
using Measurements
using LsqFit
using Latexify
using UnitfulLatexify

# Set some plotting defaults for later
# Unfortunately it doesn't accept unitformat here
default(legend=:topleft, minorgrid=:true)

# Load in data
df = DataFrame(CSV.File("dataResistance.csv", header=2, skipto=3))

# Add appropriate units to dataframe columns
df[!, "f(kHz)"] = Quantity.(df[!, "f(kHz)"], u"kHz")
df[!, 2:8] = Quantity.(df[!, 2:8], u"Ω")

# Frequency and η = ratio of resistances Rac/Rdc
freq = df[:,1]
η(col) = df[:, col] ./ df[1, col]

# Plot η against freq
plt = plot(ylabel="\\eta", unitformat=latexify)
for col in 2:8
    plot!(freq, η(col), xlabel="f", label=names(df)[col], unitformat=latexroundunitlabel)
end
display(plt)

# Plot η against √freq
sqrtPlt = plot(ylabel="\\eta", unitformat=latexify)
for col in 2:8
    plot!(sqrt.(freq), η(col), xlabel="f^{\\frac{1}{2}}", label=names(df)[col], unitformat=latexroundunitlabel)
end
display(sqrtPlt)

# Straight line model function with unknown parameters p
# It would be nested into the fit function, but it is also used in FitPlot
m(x, p) = p[1] .* x .+ p[2]

"""
Fit straight line to straight section of log-log plot. Takes a df column number for input.
"""
function fit(col::Integer, stripbelow::Real=0.1)
    ydata = log.(η(col))
    # Now we want to strip out the inital bit
    i = findfirst(x -> x >= stripbelow, ydata)
    if i === nothing
        i = 1
    end
    ydata = ydata[i:end]
    xdata = log.(freq[i:end] ./ 1u"Hz")
    p0 = [0.5, 8] # Initial parameter guess
    fitResults = curve_fit(m, xdata, ydata, p0)
    return fitResults.param
end

@userplot FitPlot
@recipe function f(o::FitPlot)
    params = o.args
    x = LinRange(0, 14, 14 * 4)
    y = m.(x, params)
    i = findfirst(x -> x > -0.1, y)
    y = y[i:end]
    x = x[i:end]

    label := false
    linestyle := :dash

    x, y
end

xintercept(p) = - p[2] / p[1]

dfAnalysed = DataFrame(Material=names(df)[2:end], fc=0.0u"Hz")

# Plot ln(η) against ln(freq)
logPlt = plot(ylabel="ln(η)")
for col in 2:8
    plot!(log.(freq ./ 1u"Hz"), log.(η(col)), xlabel="ln(f/Hz)", label=names(df)[col], unitformat=latexroundunitlabel) # Plot data
    color = logPlt[1][end][:linecolor] # Grab colour of latest line
    params = fit(col) # Lst sq fit to straight section
    xint  = xintercept(params) # Find x-intercept
    display(xint)
    fc = exp(xint) * u"Hz" # Find critical frequency from x-intercept
    dfAnalysed[col - 1, 2] = fc # Store critical frequency
    # fcPretty = round(xintercept(params); digits=2) # Prettify it
    # logPlt[1][end][:label] = "$(logPlt[1][end][:label]) fc = $fcPretty"
    fitplot!(params, linecolor=color) # Plot fitted line in same colour as source data
end
display(logPlt)

# Load in diameters
diametersRaw = DataFrame(CSV.File("dataDiameters.csv"))
avgDiam = Array(diametersRaw[end - 1, 2:end - 2]) * 1u"mm"
errDiam = Array(diametersRaw[end, 2:end - 2]) * 1u"mm"
diams = avgDiam .± errDiam
# Account for μ metal sample being rectangular, not cylindrical
μWidth = (diametersRaw[end - 1,end - 1] * 1u"mm") ± (diametersRaw[end,end - 1] * 1u"mm")
μHeight = (diametersRaw[end - 1,end] * 1u"mm") ± (diametersRaw[end,end] * 1u"mm")
μDiameter = 2 * μWidth * μHeight / (μWidth + μHeight)
push!(diams, μDiameter)
# Add to dataframe
dfAnalysed[!, :Radius] = diams ./ 2

# Conductivity, from online
σ_copper = 58.7e6u"S/m"
σ_steel = 10.1e6u"S/m"
σ = [fill(σ_copper, 4); fill(σ_steel, 2); 1/(60u"μΩ*cm")]
dfAnalysed[!, :σ] = σ

# Find μ
μ(fc, r, σ) = uconvert(u"H/m", 4 / (π * r^2 * σ * fc))
μRelative = μ.(dfAnalysed[:, :fc][1:end], dfAnalysed[:, :Radius][1:end], dfAnalysed[:, :σ][1:end]) / u"μ0"
dfAnalysed[!, :μRel] = μRelative

# Add η @ f=1MHz as column
dfAnalysed[!, :η1MHz] = Array(df[13, 2:end]) ./ Array(df[1, 2:end])
copper = sort(dfAnalysed[1:4, :], :Radius)
steel = sort(dfAnalysed[5:6, :], :Radius)

radiusPlt = plot(copper[!, :Radius], copper[!, :η1MHz], label="Copper", xlabel="Radius", ylabel="η")
plot!(steel[!, :Radius], steel[!, :η1MHz], label="Steel")
display(radiusPlt)


#========================================================================

	Steady state solution for an Aiyagari model with a diffusion process

	Translated julia code based on code from Ben Moll's Website:
		http://www.princeton.edu/%7Emoll/HACTproject.htm

		For Julia 1.0.0
========================================================================#
using LinearAlgebra, SparseArrays, Plots

# Model parameters
γ = 2 # CRRA utility
α = 1/3 # parameter for cobb-douglas production function
δ = 0.05 # captal depreciation
ρ = 0.05 # discount rate
A_prod =0.1 # term in production function

# Parameters for productivity and the poisson process q
z_1 = 1
z_2 = 2*z_1
z = [z_1 z_2]
λ_1 = 1/3 # prob of transitioning from state 1 to 2
λ_2 = 1/3  # state 2 to 1
λ = [λ_1 λ_2]
z_ave = (z_1*λ_2 + z_2*λ_1)/(λ_1 + λ_2) # Average value of z

T =200 #timing setup for shock
N = 400
dt = T/N
time = (0:N-1)*dt
max_price_it = 300
con_crit=10^(-5)
relax =0.1

# Construct the TFP sequence
corr =0.8
ν=1-corr
A_prod_t = zeros(N,1)
A_prod_t[1] = .97*A_prod

for n =1:N-1
	A_prod_t[n+1] = dt*ν*(A_prod-A_prod_t[n]) + A_prod_t[n]
end

plot(time, A_prod_t, xlims=(0,40))

# Now to set up the other grid spaces
H = 1000 # Number of elements in grid space for a
a_min = -0.8
a_max = 20.0

# Grid spaces for variables
a = LinRange(a_min, a_max, H)
a = convert(Array,a)
da = (a_max-a_min)/(H-1)

aa = [a a]
zz = ones(H,1)*z #check later

# Paramters for the simulation
maxit = 100
Ir = 40
crit = 10^(-6)
critS = 10^(-5)
Δ = 1000 #step size for HJB

# Spaces for the finite difference terms
Vaf, Vab, Vzf, Vzb, Vzz, c= [zeros(H,2) for i in 1:6]
V_r = zeros(H,2,H)

# Spaces for other terms
r_r =[]
r_min_r =[]
r_max_r =[]
KD=[]
KS=[]
S=[]
adot=[]
V_n=[]
g_r=[]
gg_s=[]

# Construct the matrix to summarize the evolution of z over time
# this matrix will not change over the finite differences and is based on the poisson process

Aswitch = [-sparse(I,H,H)*λ[1] sparse(I,H,H)*λ[2] ; sparse(I,H,H)*λ[1] -sparse(I,H,H)*λ[2]]

# Prices
r =0.04 #initial r
w = 0.05

r0=0.03
global r_min=0.01
global r_max=.99*ρ

dist = []
v0=zeros(H,2)

v0[:,1] = (w*z[1].+r.*a).^(1-γ)/(1-γ)/ρ
v0[:,2] = (w*z[2].+r.*a).^(1-γ)/(1-γ)/ρ

# The main loop, this makes sure markets clear

for ir in 1:Ir
	global r =r
	global r_min = r_min
	global r_max = r_max

	println("Main Loop iteration $(ir)")

	push!(r_r, r)
	push!(r_min_r, r_min)
	push!(r_max_r, r_max )

	push!(KD, (α*A_prod/(r+δ))^(1/(1-α))*z_ave)
	global w = (1-α)*A_prod*KD[ir].^α * z_ave^(-α)

	if ir >1
		global v0 = V_r[:,:,ir-1]
	end

	global v=v0

# Inner loop, value function iteration

	for n in 1:maxit
		global V=v
		push!(V_n, V)
		# Forward difference
		Vaf[1:H-1,:] = (V[2:H,:]-V[1:H-1,:])/da
		Vaf[H,:] = (w*z .+ r.*a_max).^(-γ)

		# Backward difference
		Vab[2:H,:] = (V[2:H,:]-V[1:H-1,:])/da
		Vab[1,:] = (w*z .+ r.*a_min).^(-γ)

		# caculate consumption and savings
		# First the forward difference case
		cf = Vaf.^(-1/γ)
		sf = w*zz + r.*aa -cf

		# consumption and savings for the backward difference
		cb = Vab.^(-1/γ)
		sb = w*zz + r.*aa -cb

		#Consumption and V' at the steady state
		c0 = w*zz + r.*aa
		Va0 = c0.^(-γ)

		# Setup to the Upwind scheme
		If = sf .> 0
		Ib = sb .< 0
		I0 = (1 .- If .- Ib)

		#Now implement the upwind scheme
		Va_Upwind = Vaf.*If + Vab.*Ib + Va0.*I0

		global c= Va_Upwind.^(-1/γ)
		u = c.^(1-γ)/(1-γ)

		X = -min.(sb,0)/da
		Y = -max.(sf,0)/da + min.(sb,0)/da
		Z = max.(sf,0)/da

		A1 = spdiagm(0=>Y[:,1], -1 => X[2:H,1], 1=> Z[1:H-1,1])
		A2 = spdiagm(0=>Y[:,2], -1 => X[2:H,2], 1=> Z[1:H-1,2])
		AA = [A1 spzeros(H,H); spzeros(H,H) A2]

		global A= AA + Aswitch

		B = (1/Δ + ρ)*sparse(I,2*H,2*H)-A

		u_stacked = [u[:,1]; u[:,2]]
		V_stacked = [V[:,1]; V[:,2]]

		b = u_stacked + V_stacked/Δ

		V_stacked = B\b

		global V = [V_stacked[1:H] V_stacked[H+1:2*H]]

		V_change = V-v

		global v = V

		push!(dist,findmax(abs.(V_change))[1])

		if n>1 && dist[n] < crit
			println("Value Function Converge Iteration = $(n)")
			break
		end
	end
	# Now solve the Fokker-Planck equation
	AT = A'
	b = zeros(2*H,1)

	# Dirty fix to avoid matrix singularities
	i_fix=1
	b[i_fix] =.1
	row = [zeros(1,i_fix-1) 1 zeros(1,H*2-i_fix)]
	AT[i_fix,:] = row

	# Now solve the linear system
	gg = AT\b
	g_sum = gg'*ones(2*H,1).*da
	gg = gg./g_sum

	global g = [gg[1:H] gg[H+1:2*H]]

	push!(gg_s,gg)
	push!(g_r, g)
	push!(adot, w*zz+r.*aa-c)
	V_r[:,:,ir] = V

	push!(KS, g[:,1]'*a*da + g[:,2]'*a*da)
	push!(S, KS[ir]-KD[ir])

	# Update aggregate capital
	if S[ir]>critS
		println("Excess Supply")
		global r_max =r
		global r = 0.5*(r+r_min)
	elseif S[ir] < -critS
		println("Excess Demand")
		global r_min =r
		global r = 0.5*(r+r_max)
	elseif abs(S[ir]) < critS
		println("Equilibrium Found!!!")
		break
	end
end

# Save some objects
v_st = v
gg_st = gg_s[end]
K_st = KS[end]
w_st =w
r_st =r
C_st = gg_s[end]'*reshape(c,H*2,1).*da
A_prod_st = A_prod

# Transition Dynamics
gg0 = gg_st

# Set up locations for preallocting values
gg_br =[]
c_t=[]
A_t=[]
C_t=[]
K_t, K_out, r_t, w_t = [zeros(N,1) for i in 1:4]
dist_it=[]

K_t = K_st*ones(N,1)

for it in 1:max_price_it
	println("Price Iteration")
	println(it)

	global w_t = (1-α)*A_prod_t.*K_t.^α .*z_ave^(-α)
	global r_t = (α*A_prod_t.*K_t.^(α-1) .*z_ave^(1-α)) .- δ

	global V=v_st

	for n in N:-1:1
		# Forward difference
		Vaf[1:H-1,:] = (V[2:H,:]-V[1:H-1,:])/da
		Vaf[H,:] = (w_t[n]*z .+ r_t[n].*a_max).^(-γ)

		# Backward difference
		Vab[2:H,:] = (V[2:H,:]-V[1:H-1,:])/da
		Vab[1,:] = (w_t[n]*z .+ r_t[n].*a_min).^(-γ)

		# caculate consumption and savings
		# First the forward difference case
		cf = Vaf.^(-1/γ)
		sf = w_t[n]*zz + r_t[n].*aa -cf

		# consumption and savings for the backward difference
		cb = Vab.^(-1/γ)
		sb = w_t[n]*zz + r_t[n].*aa -cb

		#Consumption and V' at the steady state
		c0 = w_t[n]*zz + r_t[n].*aa
		Va0 = c0.^(-γ)

		# Setup to the Upwind scheme
		If = sf .> 0
		Ib = sb .< 0
		I0 = (1 .- If .- Ib)

		#Now implement the upwind scheme
		Va_Upwind = Vaf.*If + Vab.*Ib + Va0.*I0

		global c= Va_Upwind.^(-1/γ)
		u = c.^(1-γ)/(1-γ)
		push!(c_t,c)

		X = -min.(sb,0)/da
		Y = -max.(sf,0)/da + min.(sb,0)/da
		Z = max.(sf,0)/da

		A1 = spdiagm(0=>Y[:,1], -1 => X[2:H,1], 1=> Z[1:H-1,1])
		A2 = spdiagm(0=>Y[:,2], -1 => X[2:H,2], 1=> Z[1:H-1,2])
		AA = [A1 spzeros(H,H); spzeros(H,H) A2]

		global A= AA + Aswitch

		push!(A_t,A)

		B = (1/dt + ρ)*sparse(I,2*H,2*H)-A

		u_stacked = [u[:,1]; u[:,2]]
		V_stacked = [V[:,1]; V[:,2]]

		b = u_stacked + V_stacked/dt

		V_stacked = B\b

		global V = [V_stacked[1:H] V_stacked[H+1:2*H]]
	end
	# Now update the distribution
	push!(gg_br,gg0)
	for n = 1:N
		AT = A_t[end-(n-1)]'
		push!(gg_br,(sparse(I,2*H,2*H)-AT*dt)\gg_br[n])
		global K_out[n] = (gg_br[n][1:H])'*a*da + gg_br[n][H+1:2*H]'*a*da
		push!(C_t,gg_br[n]'*reshape(c_t[n],2*H,1).*da)
	end

	push!(dist_it, findmax(abs.(K_out-K_t))[1])

	global K_t = relax.*K_out + (1-relax).*K_t

	if dist_it[it]<con_crit
		println("Equilibrium Found")
		break
	end

	global gg_br=[]
end

# Inequality measurements

Wealth_neg_t,Wealth_Gini_t,CapInc_Gini_t,Inc_Gini_t,top_inc_t,top_wealth_t,neg_frac_t = [zeros(1,400) for i in 1:7]

for n in 1:N
	Wealth_neg_t[n] = gg_br[n][1:H]'*min.(a,0).*da + gg_br[n][H+1:2*H]'*min.(a,0).*da
	g_a_count = gg_br[n][1:H] + gg_br[n][H+1:2*H]

	# Discrete Gini for checking
	g_a = g_a_count*da
	S_a = cumsum(g_a.*a,dims=1)/sum(g_a.*a)
	trapez_a = .5 .* (S_a[1]*g_a[1]+sum((S_a[2:H]+S_a[1:H-1]).*g_a[2:H]))
	Wealth_Gini_t[n] = 1 .-2*trapez_a

	# Gini of Capital income
	yk = r_t[n].*a
	g_yk=g_a

	S_yk = cumsum(g_yk.*yk,dims=1)/sum(g_yk.*yk)
	trapez_yk = .5 .* (S_yk[1]*g_yk[1]+sum((S_yk[2:H]+S_yk[1:H-1]).*g_yk[2:H]))
	CapInc_Gini_t[n] = 1 -2*trapez_yk

	# Gini of total income
	y = w_t[n]*zz + r_t[n]*aa
	Ny = 2*H
	yy =reshape(y,Ny,1)
	index = sortperm(yy[:])
	yy=yy[index]
	g_y = gg_br[n][index].*da

	S_y = cumsum(g_y.*yy,dims=1)/sum(g_y.*yy)
	trapez_y = .5 .* (S_y[1]*g_y[1]+sum((S_y[1:H-1]+S_y[1:H-1]).*g_y[2:H]))
	Inc_Gini_t[n] = 1 -2*trapez_y

	G_y =cumsum(g_y,dims=2)
	G_a = cumsum(g_a,dims=2)

	#Top 10% income share
	p1=0.1
	obj, index = findmin(abs.((1 .-G_y) .-p1))
	top_inc_t[n]=1 .-S_y[index]

	#Top 10% Wealth share
	obj, index = findmin(abs.((1 .-G_a) .-p1))
	top_wealth_t[n]=1 .-S_a[index]

	neg_frac_t[n]=findmax(G_a[a.<=0])[1]

end

p1 = plot(time, A_prod_t,xlims=(0,100), title="Aggregate Prod.",legend=false)
plot!(time,A_prod_st*ones(N,1),line=:dash,color=:black)

p2 = plot(time, K_t,xlims=(0,100),ylims=(0.293,0.3),title="Aggregate Capital St.",legend=false )
plot!(p2,time,K_st.*ones(N,1),line=:dash,color=:black)

p3 = plot(time,w_t, xlims=(0,100),title="Wage",legend=false)
plot!(p3,time,w_st.*ones(N,1),line=:dash,color=:black)

p4=plot(time,r_t,xlims=(0,100),title="Interest Rate",legend=false)
plot!(p4,time,r_st.*ones(N,1),line=:dash,color=:black)

p5 =plot(time,Wealth_Gini_t',xlims=(0,100),title="Wealth Gini",legend=false)
plot!(p5,time,Wealth_Gini_t[N].*ones(N,1),line=:dash,color=:black)

p6 =plot(time,Inc_Gini_t',xlims=(0,100),title="Income Gini",legend=false)
plot!(p6,time,Inc_Gini_t[N].*ones(N,1),line=:dash,color=:black)

plot(p1,p2,p3,p4,p5,p6,layout=(2,3),legend=false)
png("Inequality_Measures")


@doc raw"""
`SPSDE`: Stratonovich Split Partitioned Stochastic Differential Equation

Defines a partitioned stochastic differential initial value problem
```math
\begin{aligned}
dq (t) &=   v(t, q(t)) \, dt + B(t, q(t)) \circ dW , & q(t_{0}) &= q_{0} , \\
dp (t) &= [ f_1(t, q(t)) + f_2(t, q(t)) ] \, dt + [ G_1(t, q(t)) + G_2(t, q(t)) ] \circ dW , & p(t_{0}) &= p_{0} ,
\end{aligned}
```
with the drift vector fields ``v`` and ``f_i``, diffusion matrices ``B`` and ``G_i``,
initial conditions ``q_{0}`` and ``p_{0}``, the dynamical variables ``(q,p)`` taking
values in ``\mathbb{R}^{d} \times \mathbb{R}^{d}``, and the m-dimensional Wiener process W

### Parameters

* `DT <: Number`: data type
* `TT <: Real`: time step type
* `AT <: AbstractArray{DT}`: array type
* `vType <: Function`: type of `v`
* `f1Type <: Function`: type of `f1`
* `f2Type <: Function`: type of `f2`
* `BType <: Function`: type of `B`
* `G1Type <: Function`: type of `G1`
* `G2Type <: Function`: type of `G2`
* `pType <: Union{NamedTuple,Nothing}`: parameters type

### Fields

* `d`:  dimension of dynamical variable ``q`` and the vector fields ``vi``
* `m`:  dimension of the Wiener process
* `ni`: number of initial conditions
* `ns`: number of sample paths
* `v` :  function computing the drift vector field for the position variable ``q``
* `f1`:  function computing the drift vector field for the momentum variable ``p``
* `f2`:  function computing the drift vector field for the momentum variable ``p``
* `B` :  function computing the d x m diffusion matrix for the position variable ``q``
* `G1`:  function computing the d x m diffusion matrix for the momentum variable ``p``
* `G2`:  function computing the d x m diffusion matrix for the momentum variable ``p``
* `t₀`: initial time
* `q₀`: initial condition for dynamical variable ``q`` (may be a random variable itself)
* `p₀`: initial condition for dynamical variable ``p`` (may be a random variable itself)
* `parameters`: either a `NamedTuple` containing the equations parameters or `nothing`
* `periodicity`: determines the periodicity of the state vector `q` for cutting periodic solutions

The functions `v`, `f`, `B` and `G`, providing the drift vector fields and diffusion matrices, take four arguments,
`v(t, q, p, v)`, `f(t, q, p, f)`, `B(t, q, p,  B)` and `G(t, q, p, G)`, where `t` is the current time, `(q, p)` is the
current solution vector, and `v`, `f`, `B` and `G` are the variables which hold the result
of evaluating the vector fields ``v``, ``f`` and the matrices ``B``, ``G`` on `t` and `(q,p)`.

### Constructors

```julia
SPSDE(m, ns, v, f1, f2, B, G1, G2, t₀, q₀, p₀; parameters=NullParameters(), periodicity=zero(q₀[begin]))
SPSDE(m, ns, v, f1, f2, B, G1, G2, q₀::StateVector, p₀::StateVector; kwargs...) = SPSDE(m, ns, v, f1, f2, B, G1, G2, 0.0, q₀, p₀; kwargs...)
SPSDE(m, ns, v, f1, f2, B, G1, G2, t₀, q₀::State, p₀::State; kwargs...) = SPSDE(m, ns, v, f1, f2, B, G1, G2, t₀, [q₀], [p₀]; kwargs...)
SPSDE(m, ns, v, f1, f2, B, G1, G2, q₀::State, p₀::State; kwargs...) = SPSDE(m, ns, v, f1, f2, B, G1, G2, 0.0, q₀, p₀; kwargs...)
```

### Example

```julia
    function v(λ, t, q, v)
        v[1] = λ*q[1]
        v[2] = λ*q[2]
    end

    function B(μ, t, q, B)
        B[1] = μ*q[1]
        B[2] = μ*q[2]
    end

    t₀ = 0.
    q₀ = [1., 1.]
    λ  = 2.
    μ  = 1.

    v_sde = (t, q, v) -> v(λ, t, q, v)
    B_sde = (t, q, B) -> B(μ, t, q, B)

    sde = SDE(v_sde, B_sde, t₀, q₀)
```
"""
struct SPSDE{dType <: Number, tType <: Real, arrayType <: AbstractArray{dType},
             vType <: Function, f1Type <: Function, f2Type <: Function,
             BType <: Function, G1Type <: Function, G2Type <: Function,
             pType <: OptionalParameters} <: AbstractEquationPSDE{Nothing,Nothing,Nothing}
    d::Int
    m::Int
    ns::Int
    v ::vType
    f1::f1Type
    f2::f2Type
    B ::BType
    G1::G1Type
    G2::G2Type
    t₀::tType
    q₀::Vector{arrayType}
    p₀::Vector{arrayType}
    parameters::pType
    periodicity::arrayType

    function SPSDE(m, ns,
                    v::vType, f1::f1Type, f2::f2Type,
                    B::BType, G1::G1Type, G2::G2Type,
                    t₀::tType, q₀::Vector{arrayType}, p₀::Vector{arrayType};
                    parameters=NullParameters(), periodicity=zero(q₀[begin])) where {
                        dType <: Number, tType <: Real, arrayType <: AbstractArray{dType},
                        vType <: Function, f1Type <: Function, f2Type <: Function,
                        BType <: Function, G1Type <: Function, G2Type <: Function}

        d  = length(q₀[begin])
        ni = length(q₀)

        @assert length(q₀) == length(p₀)
        @assert all(length(q) == d for q in q₀)
        @assert all(length(p) == d for p in p₀)

        @assert ns ≥ 1
        @assert ni == 1 || ns == 1
        # either multiple deterministic initial conditions and one sample path
        # or one deterministic initial condition and multiple sample paths

        new{dType,tType,arrayType,vType,f1Type,f2Type,BType,G1Type,G2Type,typeof(parameters)}(d, m, ns, v, f1, f2, B, G1, G2, t₀, q₀, p₀, parameters, periodicity)
    end
end

SPSDE(m, ns, v, f1, f2, B, G1, G2, q₀::StateVector, p₀::StateVector; kwargs...) = SPSDE(m, ns, v, f1, f2, B, G1, G2, 0.0, q₀, p₀; kwargs...)
SPSDE(m, ns, v, f1, f2, B, G1, G2, t₀, q₀::State, p₀::State; kwargs...) = SPSDE(m, ns, v, f1, f2, B, G1, G2, t₀, [q₀], [p₀]; kwargs...)
SPSDE(m, ns, v, f1, f2, B, G1, G2, q₀::State, p₀::State; kwargs...) = SPSDE(m, ns, v, f1, f2, B, G1, G2, 0.0, q₀, p₀; kwargs...)

const SPSDEPT{PT,DT,TT,AT,VT,F1T,F2T,BT,G1T,G2T} = SPSDE{DT,TT,AT,VT,F1T,F2T,BT,G1T,G2T,PT} # type alias for dispatch on parameters type parameter

Base.hash(sde::SPSDE, h::UInt) = hash(sde.d, hash(sde.m, hash(sde.ns,
                                 hash(sde.v, hash(sde.f1, hash(sde.f2,
                                 hash(sde.B, hash(sde.G1, hash(sde.G2,
                                 hash(sde.t₀, hash(sde.q₀, hash(sde.p₀,
                                 hash(sde.parameters, hash(sde.periodicity, h))))))))))))))

Base.:(==)(sde1::SPSDE, sde2::SPSDE) = (
                                sde1.d == sde2.d
                             && sde1.m == sde2.m
                             && sde1.ns == sde2.ns
                             && sde1.v  == sde2.v
                             && sde1.f1 == sde2.f1
                             && sde1.f2 == sde2.f2
                             && sde1.B  == sde2.B
                             && sde1.G1 == sde2.G1
                             && sde1.G2 == sde2.G2
                             && sde1.t₀ == sde2.t₀
                             && sde1.q₀ == sde2.q₀
                             && sde1.p₀ == sde2.p₀
                             && sde1.parameters  == sde2.parameters
                             && sde1.periodicity == sde2.periodicity)

function Base.similar(equ::SPSDE, t₀::Real, q₀::StateVector, p₀::StateVector, ns::Int=equ.ns)
    @assert all([length(q) == ndims(equ) for q in q₀])
    @assert all([length(p) == ndims(equ) for p in p₀])
    SPSDE(equ.m, ns, equ.v, equ.f1, equ.f2, equ.B, equ.G1, equ.G2, t₀, q₀, p₀; parameters=equ.parameters, periodicity=equ.periodicity)
end

Base.similar(equ::SPSDE, t₀::Real, q₀::State, p₀::State, ns::Int=equ.ns; kwargs...) = similar(equ, t₀, [q₀], [p₀], ns; kwargs...)
Base.similar(equ::SPSDE, q₀::Union{State,StateVector}, p₀::Union{State,StateVector}, ns::Int=equ.ns; kwargs...) = similar(equ, equ.t₀, q₀, p₀, ns; kwargs...)

Base.ndims(equ::SPSDE) = equ.d
Base.axes(equ::SPSDE) = axes(equ.q₀[begin])
GeometricBase.nsamples(equ::SPSDE) = length(equ.q₀)
GeometricBase.periodicity(equ::SPSDE) = equ.periodicity

initial_conditions(equ::SPSDE) = (equ.t₀, equ.q₀, equ.p₀)

hasparameters(::SPSDEPT{<:NullParameters}) = false
hasparameters(::SPSDEPT{<:NamedTuple}) = true

_get_v( equ::SPSDE) = hasparameters(equ) ? (t,q,p,v) -> equ.v( t, q, p, v, equ.parameters) : equ.v
_get_f1(equ::SPSDE) = hasparameters(equ) ? (t,q,p,f) -> equ.f1(t, q, p, f, equ.parameters) : equ.f1
_get_f2(equ::SPSDE) = hasparameters(equ) ? (t,q,p,f) -> equ.f2(t, q, p, f, equ.parameters) : equ.f2
_get_B( equ::SPSDE) = hasparameters(equ) ? (t,q,p,B) -> equ.B( t, q, p, B, equ.parameters) : equ.B
_get_G1(equ::SPSDE) = hasparameters(equ) ? (t,q,p,G) -> equ.G1(t, q, p, G, equ.parameters) : equ.G1
_get_G2(equ::SPSDE) = hasparameters(equ) ? (t,q,p,G) -> equ.G2(t, q, p, G, equ.parameters) : equ.G2

function functions(equ::SPSDE)
    names = (:v,:f1,:f2,:B,:G1,:G2)
    equs  = (_get_v(equ), _get_f1(equ), _get_f2(equ), _get_B(equ), _get_G1(equ), _get_G2(equ))

    NamedTuple{names}(equs)
end


function poly_horner(x, a...)
b=zero(x)
for i = length(a):-1:1
b = a[i] + b * x
end
return b
end


using LinearAlgebra
using MPOmodule

sz = Array{ComplexF64}([1.0 0.0; 0.0 -1.0])
sy = Array{ComplexF64}([0.0 -1.0im; 1.0im 0.0])
sx = Array{ComplexF64}([0.0 1.0; 1.0 0.0])
Id(n) = Array{ComplexF64}(I,n,n)
Basis = [Id(2), sx, sy, sz]

ca = Array{ComplexF64}([0.0 1.0; 0.0 0.0])
cr = Array{ComplexF64}([0.0 0.0; 1.0 0.0])
#n = Array{ComplexF64}([1.0 0.0; 0.0 0.0]) #changed this
n = Array{ComplexF64}([0.0 0.0; 0.0 1.0]) #changed this

#=========================================================

		MPO models

=========================================================#

""" SZ - measurement """
Op_vec_Id = Vector{localOp}(localOp.([Id(2),Id(2),Id(2)]))
Op_vec_Sz = Vector{localOp}(localOp.([Id(2),sz,Id(2)]))
Op_idx_Sz = Vector{Int}([1,2,4])

SzMPO = MPO(2, Op_vec_Sz, Op_idx_Sz)

"""occupation"""

#
#	1	0
#	n	1
#

Op_vec_n = Vector{localOp}(localOp.([Id(2),n,Id(2)]))
Op_idx_n = Vector{Int}([1,2,4])

Op_vec_n_start = Vector{localOp}(localOp.([n, Id(2)]))
Op_idx_n_start = Vector{Int}([1,2])

Op_vec_n_end = Vector{localOp}(localOp.([Id(2), n]))
Op_idx_n_end = Vector{Int}([1,2])

""" transver Ising """

# creat MPO representation for transver ising
#
#	 1  	0	0
#	-J*sx	0	0
#	-h*sz	sx	1
#

Op_vec_TIsing(J, h) = Vector{localOp}(localOp.([Id(2), -J*sx, -h*sz, sx, Id(2)]))
Op_idx_TIsing = Vector{Int}([1, 2, 3, 6, 9])

Op_vec_TIsing_start(J, h) = Vector{localOp}(localOp.([-h*sz, sx, Id(2)]))
Op_idx_TIsing_start = Vector{Int}([1, 2, 3])

Op_vec_TIsing_end(J, h) = Vector{localOp}(localOp.([Id(2), -J*sx, -h*sz]))
Op_idx_TIsing_end = Vector{Int}([1, 2, 3])

""" Rice - Mele with intra cell repulsion """

#
#
#	1		 0	  0	0	0
#	c^†		 0	  0	0	0
#	c		 0	  0	0	0
#	n		 0	  0	0	0
#	-(V - μ)n 	-Jc	-Jc^†	Un	1
#

Op_vec_RM(J, V, U, μ) = Vector{localOp}(localOp.([Id(2), cr, ca, n, -(V-μ)*n, -J*ca, -J*cr, U*n, Id(2)]))
Op_idx_RM = Vector{Int}([1, 2, 3, 4, 5, 10, 15, 20, 25])

Op_vec_RM_start(J, V, U, μ) = Vector{localOp}(localOp.([-(V-μ)*n, -J*ca, -J*cr, U*n, Id(2)]))
Op_idx_RM_start = Vector{Int}([1,2,3,4,5])

Op_vec_RM_end(V, μ) = Vector{localOp}(localOp.([Id(2), cr, ca, n, -(V-μ)*n]))
Op_idx_RM_end = Vector{Int}([1,2,3,4,5])

""" Cij matrix """

function CMPO(i,j,N)

	Id_local = localOp(Id(2))
	CijVec = [MPO(1, Vector{localOp}([Id_local]), Vector{Int}([1])) for k in 1:N]
	CijVec[i] = MPO(1, localOp.([cr]), Vector{Int}([1]))
	#CijVec[i] = MPO(1,[ca],[1])
	if i!=j
		CijVec[j] = MPO(1, localOp.([ca]), Vector{Int}([1]))
		#CijVec[j] = MPO(1,[cr],[1])
	else
		CijVec[i] = MPO(1,  localOp.([cr*ca]), Vector{Int}([1]))
		#CijVec[i] = MPO(1,[ca*cr],[1])
	end

	for k in Int(i+1):Int(j-1)
		CijVec[k] = MPO(1, localOp.([(Id(2)-2*cr*ca)]), Vector{Int}([1]))#[(Id(2)-2*cr*ca)], [1])
		#CijVec[k] = MPO(1, [(Id(2)-2*ca*cr)], [1])#[(Id(2)-2*cr*ca)], [1])
	end

	return Vector{MPO}(CijVec)
end


vars=5
i_train=10_000
i_test=10_000
cc1=10.0
neurons=Int64(floor(i_train/50)) # this automatically
# generates number of neurons== number of clusters (§3.1)
_case_noise_=5

ff(x)=-x[1]+(1/2)x[2].^2-(1/3)x[3].^3+(1/4)x[4].^4-(1/5)x[5].^5
xx_train=0.8*rand(i_train,vars).+0.1
yy_train=zeros(i_train); for i=1:i_train yy_train[i]=ff(xx_train[i,:]) end
mi=minimum(yy_train);yy_train.-=mi;ma=maximum(yy_train);yy_train./=ma;yy_train.*=0.8;yy_train.+=0.1

if _case_noise_==1
    __noise__=rand(i_train)
elseif _case_noise_==2
    __noise__=rand(Normal(0.0, 1.0),i_train)
elseif _case_noise_==3
    __noise__=rand(GeneralizedPareto(0.0,1.0,0.25),i_train)
elseif _case_noise_==4
    __noise__=rand(LogNormal(0.25, 0.8),i_train)
elseif _case_noise_==5
    __noise__=[rand(LogNormal(0.25, 0.8),i_train);10.0.-rand(Exponential(3.0),i_train)]
    __noise__=[__noise__;15.0.+rand(Frechet(10.0,10.0),i_train)]
    __noise__=__noise__[1:3:end]
end
__noise__.-=mean(__noise__)
__noise__./=mean(abs.(__noise__))/(0.05*0.45)
histogram(__noise__,label=sum(__noise__)/sum(abs.(__noise__)),
        title=mean(abs.(__noise__)))
iso=sortperm(yy_train)
plot(yy_train[iso])
yy_train.+=__noise__
display(plot!(yy_train[iso],legend=false))

xx_test=0.8*rand(i_test,vars).+0.1
yy_test=zeros(i_test); for i=1:i_test yy_test[i]=ff(xx_test[i,:]) end
# yy_test.+=noise_perc.*rand(i_test)
yy_test.-=mi;yy_test./=ma;yy_test.*=0.8;yy_test.+=0.1

scatter3d(xx_train[:,1],xx_train[:,2],yy_train)
display(scatter3d!(xx_test[:,1],xx_test[:,2],yy_test))

using DecisionTree, Statistics
num_trees=100;
regr_3 = RandomForestRegressor(n_trees=num_trees);
tmp_xx_train=copy(xx_train);tmp_yy_train=copy(yy_train);
tmp_xx_test=copy(xx_test);tmp_yy_test=copy(yy_test);
DecisionTree.fit!(regr_3, tmp_xx_train, tmp_yy_train[:,1]);
predl1 = DecisionTree.predict(regr_3, tmp_xx_train);
predl1_test = DecisionTree.predict(regr_3, tmp_xx_test);
mean(abs.(predl1-yy_train))
mae=mean(abs.(predl1_test-yy_test))
err=yy_test-predl1_test;ier=sortperm(err)
display(plot(err[ier],label="Random Forests",color=:black,
    linestyle=:dash,legend=:topleft,linewidth=3))
__err__=string(mean(abs.(err)))

sci_model = fit!(AdaBoostRegressor(n_estimators=100),xx_train,yy_train)
predl1=ScikitLearn.predict(sci_model,xx_train)
predl1_test=ScikitLearn.predict(sci_model,xx_test)
maetr=mean(abs.(yy_train-predl1))
maete=mean(abs.(yy_test-predl1_test))
err=yy_test-predl1_test;ier=sortperm(err)
display(plot!(err[ier],label="AdaBoost",color=:black,
    linestyle=:dashdot,legend=:topleft,linewidth=3))
__err__=string(__err__," ",mean(abs.(err)))

num_round = 10
bst = xgboost(copy(xx_train), num_round, label = copy(yy_train), eta = 1, max_depth = 7)
predl1 = XGBoost.predict(bst, copy(xx_train))
predl1_test = XGBoost.predict(bst, copy(xx_test))
mean(abs.(predl1-yy_train))
mae=mean(abs.(predl1_test-yy_test))
err=yy_test-predl1_test;ier=sortperm(err)
display(plot!(err[ier],label="XGBoost",color=:black,
    linestyle=:dot,legend=:topleft,linewidth=3))
__err__=string(__err__," ",mean(abs.(err)))

##### ANNBN RBF

include(string(path1,"/src/ANNBN.jl"))
inds_all,n_per_part=ANNBN.___clustering(neurons,xx_train,200)

a_all,a_layer1,layer1=ANNBN.train_layer_1_rbf(neurons,vars,i_train,n_per_part,inds_all,xx_train,yy_train,cc1)
predl1,layer1_train=ANNBN.predict_new_rbf(a_all,a_layer1,xx_train,i_train,neurons,n_per_part,inds_all,xx_train,vars,cc1)
predl1_test,layer1_test=ANNBN.predict_new_rbf(a_all,a_layer1,xx_test,i_test,neurons,n_per_part,inds_all,xx_train,vars,cc1)
maetr=mean(abs.(yy_train-predl1))
maete=mean(abs.(yy_test-predl1_test))
err=yy_test-predl1_test;ier=sortperm(err)
display(plot!(err[ier],label="ANNBN",color=:black,linestyle=:dashdotdot,
    legend=:topleft,linewidth=3))
__err__=string(__err__," ",mean(abs.(err)))

println(__err__)

#


module OptimalTransport

# MONGE PROBLEM
# -------------

using Combinatorics: permutations

X1 = [-3.334062779243765 -0.4501346178106358;
-17.089568765347632 13.122166467239909;
-14.347484681558905 -6.733978925249632;
4.023051337264342 -9.962304692857852;
-8.463617750898711 -7.272060919414795;
4.411632356371395 18.33766410251966;
0.33231514398722767 -8.943403560908134;
4.312466564646311 -4.769277598451736;
-0.1530439721692118 9.400009756002328;
-5.9829675943508995 0.07220200778660099];

X2 = [-15.776290315139038 -5.26704453200824;
5.035133784625716 6.388525577583841;
-1.302878737641465 -14.087978411169559;
-17.05014605333986 10.189492174622178;
-10.102710138685817 1.6112603598816224;
-17.880416477734084 -5.15108311036524;
-18.49092111472347 -5.193032093476489;
1.2972905950753137 34.72586483583929;
1.3580992685186684 7.334727533023157;
-7.482909025521323 6.326317440680775];

function monge_brute_force(C)
  n, m = size(C)
  @assert n == m "C should be square"
  # loop over all permutations and to find the
  # matching with the lowest cost

  best_cost = typemax(eltype(C))
  for p in permutations(1:n)
      cost = zero(eltype(C))
      for (i, j) in enumerate(p)
        cost += C[i,j]
      end
      if cost < best_cost
        best_cost = cost
        global best_perm = p
      end
    end
  return best_perm, best_cost
end

export X1, X2, monge_brute_force

# SINKHORN
# -------

using LinearAlgebra

function sinkhorn(C::Matrix, a::Vector, b::Vector; λ=1.0, ϵ=1e-8)
    n, m = size(C)
    @assert n == length(a) && m == length(b) throw(DimensionMismatch("a and b do not match"))
    @assert sum(a) ≈ sum(b) "a and b don't have equal sums"
    u, v = copy(a), copy(b)
    M = exp.(-λ * C)
    # normalize this matrix
    while maximum(abs.(a .- Diagonal(u) * (M * v))) > ϵ
        u .= a ./ (M * v)
        v .= b ./ (M' * u)
      end
    return Diagonal(u) * M * Diagonal(v)
  end

export sinkhorn

end


module Diffusion

import MPI
using ImplicitGlobalGrid
using CUDA
using Plots

@views d_xa(A) = A[2:end  , :     ] .- A[1:end-1, :     ];
@views d_xi(A) = A[2:end  ,2:end-1] .- A[1:end-1,2:end-1];
@views d_ya(A) = A[ :     ,2:end  ] .- A[ :     ,1:end-1];
@views d_yi(A) = A[2:end-1,2:end  ] .- A[2:end-1,1:end-1];
@views  inn(A) = A[2:end-1,2:end-1]

@views function diffusion2D(nx, ny, nt, lx, ly, lam, cp_min, amp_T, Cp, comm)
    # Numerics
    me, dims = init_global_grid(nx, ny, 1; init_MPI=false, comm=comm, reorder=0)  # Initialize the implicit global grid
    dx       = lx/(nx_g()-1)                                                      # Space step in dimension x
    dy       = ly/(ny_g()-1)                                                      # ...        in dimension y

    # Array initializations
    T     = CUDA.zeros(Float64, nx,   ny, )
    dTedt = CUDA.zeros(Float64, nx-2, ny-2)
    qx    = CUDA.zeros(Float64, nx-1, ny-2)
    qy    = CUDA.zeros(Float64, nx-2, ny-1)

    # Initialize temperature with a Gaussian anomaly of amplitude amp_T
    T    .= CuArray([amp_T*exp(-((x_g(ix,dx,T)-lx/2)/2)^2-((y_g(iy,dy,T)-ly/2)/2)^2) for ix=1:size(T,1), iy=1:size(T,2)])

    # Preparation of visualisation
    gr()
    ENV["GKSwstype"]="nul"
    anim     = Animation();
    nx_v     = (nx-2)*dims[1]
    ny_v     = (ny-2)*dims[2]
    T_v      = zeros(nx_v, ny_v)
    T_nohalo = zeros(nx-2, ny-2)

    # Time loop
    nframes = 50
    dt = min(dx,dy)^2*cp_min/lam/4.1                            # Time step for the 2D Heat diffusion
    for it = 1:nt
        qx    .= -lam.*d_xi(T)./dx                              # Fourier's law of heat conduction: q_x   = -λ δT/δx
        qy    .= -lam.*d_yi(T)./dy                              # ...                               q_y   = -λ δT/δy
        dTedt .= 1.0./inn(Cp).*(-d_xa(qx)./dx .- d_ya(qy)./dy)  # Conservation of energy:           δT/δt = 1/cₚ (-δq_x/δx - δq_y/dy)
        T[2:end-1,2:end-1] .= inn(T) .+ dt.*dTedt               # Update of temperature             T_new = T_old + δT/δt
        update_halo!(T)                                         # Update the halo of T
        if mod(it, nt/nframes) == 1
            T_nohalo .= T[2:end-1,2:end-1]                      # Copy data to CPU removing the halo.
            gather!(T_nohalo, T_v)                              # Gather the global temperature field on process 0.
            if (me==0)                                          # Visualize it on process 0.
                println("Time step $it: maximum temperature is $(maximum(T_v))")
                heatmap(T_v, title="Temperature at time step $it", aspect_ratio=1, xlims=(1,nx_v), ylims=(1,ny_v), clims=(0,amp_T)); frame(anim)
            end
        end
    end

    # Postprocessing
    if (me==0) gif(anim, "diffusion2D.gif", fps = 15) end       # Create an animated gif on process 0.
    finalize_global_grid(;finalize_MPI=false)
end

Base.@ccallable function julia_diffusion(nx::Csize_t, ny::Csize_t, nt::Csize_t, lx::Cdouble, ly::Cdouble, lam::Cdouble, cp_min::Cdouble, amp_T::Cdouble, Cp_c::CuPtr{Cdouble}, comm_c::MPI.MPI_Comm)::Cint
    try
        comm = MPI.Comm_dup(MPI.Comm(comm_c))
        Cp   = unsafe_wrap(CuArray, Cp_c, (Int64(nx), Int64(ny)))
        diffusion2D(nx, ny, nt, lx, ly, lam, cp_min, amp_T, Cp, comm)
    catch
        Base.invokelatest(Base.display_error, Base.catch_stack())
        return 1
    end
    return 0
end

end # module


"""
```
measurement(m::AnSchorfheide{T}, TTT::Matrix{T},
            RRR::Matrix{T}, CCC::Vector{T}) where {T <: AbstractFloat}
```

Assign measurement equation

```
y_t = ZZ*s_t + DD + u_t
```

where

```
Var(ϵ_t) = QQ
Var(u_t) = EE
Cov(ϵ_t, u_t) = 0
```
"""
function measurement(m::AnSchorfheide{T}, TTT::Matrix{T}, # do not edit input arguments
                     RRR::Matrix{T}, CCC::Vector{T};
                     reg::Int = 1, # these keyword arguments are only used to make regime-switching work.
                     TTTs::Vector{<: AbstractMatrix{T}} = Matrix{T}[], # you can delete these keyword arguments if
                     CCCs::Vector{<: AbstractVector{T}} = Vector{T}[], # you don't need regime-switching.
                     information_set::UnitRange = reg:reg,
                     memo::Union{ForwardMultipleExpectationsMemo, Nothing} = nothing) where {T <: AbstractFloat}

    endo     = m.endogenous_states           # make references to model indices
    endo_new = m.endogenous_states_augmented # for convenience. You can rename thse if you want
    exo      = m.exogenous_shocks
    obs      = m.observables

    _n_observables = n_observables(m) # do not edit!
    _n_states = n_states_augmented(m)
    _n_shocks_exogenous = n_shocks_exogenous(m)

    ZZ = zeros(_n_observables, _n_states) # do not edit!
    DD = zeros(_n_observables)
    EE = zeros(_n_observables, _n_observables)
    QQ = zeros(_n_shocks_exogenous, _n_shocks_exogenous)

    # Starting from here, you should probably edit!

    ## Output growth
    ZZ[obs[:obs_gdp], endo[:y_t]]  = 1.0
    ZZ[obs[:obs_gdp], endo[:y_t1]] = -1.0
    ZZ[obs[:obs_gdp], endo[:z_t]]  = 1.0
    DD[obs[:obs_gdp]]              = m[:γ_Q]

    ## Inflation
    ZZ[obs[:obs_cpi], endo[:π_t]] = 4.0
    DD[obs[:obs_cpi]]             = m[:π_star]

    ## Federal Funds Rate
    ZZ[obs[:obs_nominalrate], endo[:R_t]] = 4.0
    DD[obs[:obs_nominalrate]]             = m[:π_star] + m[:rA] + 4.0*m[:γ_Q]

    # Measurement error
    EE[obs[:obs_gdp], obs[:obs_gdp]]                 = m[:e_y]^2
    EE[obs[:obs_cpi], obs[:obs_cpi]]                 = m[:e_π]^2
    EE[obs[:obs_nominalrate], obs[:obs_nominalrate]] = m[:e_R]^2

    # Variance of innovations
    QQ[exo[:z_sh],exo[:z_sh]]   = (m[:σ_z])^2
    QQ[exo[:g_sh],exo[:g_sh]]   = (m[:σ_g])^2
    QQ[exo[:rm_sh],exo[:rm_sh]] = (m[:σ_R])^2

    return Measurement(ZZ, DD, QQ, EE) # do not edit this return
end


using Statistics, Printf

using Oceananigans
using Oceananigans.Fields
using Oceananigans.TurbulenceClosures
using Oceananigans.OutputWriters
using Oceananigans.Diagnostics
using Oceananigans.Utils

""" Friction velocity. See equation (16) of Vreugdenhil & Taylor (2018). """
function uτ(model, Uavg, U_wall)
    Nz, Hz, Δz = model.grid.Nz, model.grid.Hz, model.grid.Δz
    ν = model.closure.ν

    compute!(Uavg)
    U = Array(interior(Uavg))  # Exclude average of halo region.

    # Use a finite difference to calculate dU/dz at the top and bottomtom walls.
    # The distance between the center of the cell adjacent to the wall and the
    # wall itself is Δz/2.
    uτ²_top    = ν * abs(U[1] - U_wall)    / (Δz/2)  # Top wall    where u = +U_wall
    uτ²_bottom = ν * abs(-U_wall  - U[Nz]) / (Δz/2)  # Bottom wall where u = -U_wall

    uτ_top, uτ_bottom = √uτ²_top, √uτ²_bottom

    return uτ_top, uτ_bottom
end

""" Heat flux at the wall. See equation (16) of Vreugdenhil & Taylor (2018). """
function q_wall(model, Tavg, Θ_wall)
    Nz, Hz, Δz = model.grid.Nz, model.grid.Hz, model.grid.Δz
    κ = model.closure.κ.T

    compute!(Tavg)
    Θ = Array(interior(Tavg)) # Exclude average of halo region.

    # Use a finite difference to calculate dθ/dz at the top and bottomtom walls.
    # The distance between the center of the cell adjacent to the wall and the
    # wall itself is Δz/2.
    q_wall_top    = κ * abs(Θ[1] - Θ_wall)   / (Δz/2)  # Top wall    where Θ = +Θ_wall
    q_wall_bottom = κ * abs(-Θ_wall - Θ[Nz]) / (Δz/2)  # Bottom wall where Θ = -Θ_wall

    return q_wall_top, q_wall_bottom
end

struct FrictionReynoldsNumber{H, U}
    Uavg :: H
    U_wall :: U
end

struct NusseltNumber{H, T}
    Tavg :: H
    Θ_wall :: T
end

""" Friction Reynolds number. See equation (20) of Vreugdenhil & Taylor (2018). """
function (Reτ::FrictionReynoldsNumber)(model)
    ν = model.closure.ν
    h = model.grid.Lz / 2
    uτ_top, uτ_bottom = uτ(model, Reτ.Uavg, Reτ.U_wall)

    return h * uτ_top / ν, h * uτ_bottom / ν
end

""" Nusselt number. See equation (20) of Vreugdenhil & Taylor (2018). """
function (Nu::NusseltNumber)(model)
    κ = model.closure.κ.T
    h = model.grid.Lz / 2

    q_wall_top, q_wall_bottom = q_wall(model, Nu.Tavg, Nu.Θ_wall)

    return (q_wall_top * h)/(κ * Nu.Θ_wall), (q_wall_bottom * h)/(κ * Nu.Θ_wall)
end

"""
    simulate_stratified_couette_flow(; Nxy, Nz, h=1, U_wall=1, Re=4250, Pr=0.7,
                                     Ri, Ni=10, end_time=1000)

Simulate stratified plane Couette flow with `Nxy` grid cells in each horizontal
direction, `Nz` grid cells in the vertical, in a domain of size (4πh, 2πh, 2h),
with wall velocities of `U_wall` at the top and -`U_wall` at the bottom, at a Reynolds
number `Re, Prandtl number `Pr`, and Richardson number `Ri`.

`Ni` is the number of "intermediate" time steps taken at a time before printing a progress
statement and updating the time step.
"""
function simulate_stratified_couette_flow(; Nxy, Nz, arch=GPU(), h=1, U_wall=1,
                                          Re=4250, Pr=0.7, Ri, Ni=10, end_time=1000)
    #####
    ##### Computed parameters
    #####

         ν = U_wall * h / Re    # From Re = U_wall h / ν
    Θ_wall = Ri * U_wall^2 / h  # From Ri = L Θ_wall / U_wall²
         κ = ν / Pr             # From Pr = ν / κ

    #####
    ##### Impose boundary conditions
    #####

    grid = RegularRectilinearGrid(size = (Nxy, Nxy, Nz), extent = (4π*h, 2π*h, 2h))

    Tbcs = FieldBoundaryConditions(top = BoundaryCondition(Value,  Θ_wall),
                                   bottom = BoundaryCondition(Value, -Θ_wall))

    ubcs = FieldBoundaryConditions(top = BoundaryCondition(Value,  U_wall),
                                   bottom = BoundaryCondition(Value, -U_wall))

    vbcs = FieldBoundaryConditions(top = BoundaryCondition(Value, 0),
                                   bottom = BoundaryCondition(Value, 0))

    #####
    ##### Non-dimensional model setup
    #####

    model = NonhydrostaticModel(
               architecture = arch,
                       grid = grid,
                   buoyancy = SeawaterBuoyancy(equation_of_state=LinearEquationOfState(α=1.0, β=0.0)),
                    closure = AnisotropicMinimumDissipation(ν=ν, κ=κ),
        boundary_conditions = (u=ubcs, v=vbcs, T=Tbcs)
    )

    #####
    ##### Set initial conditions
    #####

    # Add a bit of surface-concentrated noise to the initial condition
    ε(σ, z) = σ * randn() * z/model.grid.Lz * (1 + z/model.grid.Lz)

    # We add a sinusoidal initial condition to u to encourage instability.
    T₀(x, y, z) = 2Θ_wall * (1/2 + z/model.grid.Lz) * (1 + ε(5e-1, z))
    u₀(x, y, z) = 2U_wall * (1/2 + z/model.grid.Lz) * (1 + ε(5e-1, z)) * (1 + 0.5*sin(4π/model.grid.Lx * x))
    v₀(x, y, z) = ε(5e-1, z)
    w₀(x, y, z) = ε(5e-1, z)

    set!(model, u=u₀, v=v₀, w=w₀, T=T₀)

    #####
    ##### Print simulation banner
    #####

    @printf(
        """
        Simulating stratified plane Couette flow

                N : %d, %d, %d
                L : %.3g, %.3g, %.3g
               Re : %.3f
               Ri : %.3f
               Pr : %.3f
                ν : %.3g
                κ : %.3g
           U_wall : %.3f
           Θ_wall : %.3f

        """, model.grid.Nx, model.grid.Ny, model.grid.Nz,
             model.grid.Lx, model.grid.Ly, model.grid.Lz,
             Re, Ri, Pr, ν, κ, U_wall, Θ_wall)

    #####
    ##### Set up field output writer
    #####

    base_dir = @sprintf("stratified_couette_flow_data_Nxy%d_Nz%d_Ri%.2f", Nxy, Nz, Ri)
    prefix = @sprintf("stratified_couette_flow_Nxy%d_Nz%d_Ri%.2f", Nxy, Nz, Ri)

    function init_save_parameters_and_bcs(file, model)
        file["parameters/reynolds_number"] = Re
        file["parameters/richardson_number"] = Ri
        file["parameters/prandtl_number"] = Pr
        file["parameters/viscosity"] = ν
        file["parameters/diffusivity"] = κ
        file["parameters/wall_velocity"] = U_wall
        file["parameters/wall_temperature"] = Θ_wall
    end

    fields = Dict(
        :u => model -> Array(model.velocities.u.data.parent),
        :v => model -> Array(model.velocities.v.data.parent),
        :w => model -> Array(model.velocities.w.data.parent),
        :T => model -> Array(model.tracers.T.data.parent),
   :kappaT => model -> Array(model.diffusivity_fields.κₑ.T.data.parent),
       :nu => model -> Array(model.diffusivity_fields.νₑ.data.parent))

    field_writer =
        JLD2OutputWriter(model, fields, dir=base_dir, prefix=prefix * "_fields",
                         init=init_save_parameters_and_bcs, schedule=TimeInterval(10),
                         force=true, verbose=true)

    #####
    ##### Set up profile output writer
    #####

    Uavg = AveragedField(model.velocities.u,            dims=(1, 2))
    Vavg = AveragedField(model.velocities.v,            dims=(1, 2))
    Wavg = AveragedField(model.velocities.w,            dims=(1, 2))
    Tavg = AveragedField(model.tracers.T,               dims=(1, 2))
    νavg = AveragedField(model.diffusivity_fields.νₑ,   dims=(1, 2))
    κavg = AveragedField(model.diffusivity_fields.κₑ.T, dims=(1, 2))

    profiles = Dict(
         :u => Uavg,
         :v => Vavg,
         :w => Wavg,
         :T => Tavg,
        :nu => νavg,
    :kappaT => κavg)

    profile_writer =
        JLD2OutputWriter(model, profiles, dir=base_dir, prefix=prefix * "_profiles",
                         init=init_save_parameters_and_bcs, schedule=TimeInterval(1),
                         force=true, verbose=true)

    #####
    ##### Set up statistic output writer
    #####

    Reτ = FrictionReynoldsNumber(Uavg, U_wall)
     Nu = NusseltNumber(Tavg, Θ_wall)

    statistics = Dict(
        :Re_tau => model -> Reτ(model),
        :Nu     => model -> Nu(model))

    statistics_writer =
        JLD2OutputWriter(model, statistics, dir=base_dir, prefix=prefix * "_statistics",
                         init=init_save_parameters_and_bcs, schedule=TimeInterval(1/2),
                         force=true, verbose=true)

    #####
    ##### Time stepping
    #####

    wizard = TimeStepWizard(cfl=0.02, Δt=0.0001, max_change=1.1, max_Δt=0.02)

    # We will ramp up the CFL used by the adaptive time step wizard during spin up.
    cfl(t) = min(0.01t, 0.1)

    function print_progress(simulation)
        model = simulation.model
        clock = model.clock

        wizard.cfl = cfl(model.clock.time)

        progress = 100 * (clock.time / end_time)

        umax = maximum(abs, model.velocities.u.data.parent)
        vmax = maximum(abs, model.velocities.v.data.parent)
        wmax = maximum(abs, model.velocities.w.data.parent)
        CFL = wizard.Δt / cell_advection_timescale(model)

        Δ = min(model.grid.Δx, model.grid.Δy, model.grid.Δz)
        νmax = maximum(model.diffusivity_fields.νₑ.data.parent)
        κmax = maximum(model.diffusivity_fields.κₑ.T.data.parent)
        νCFL = wizard.Δt / (Δ^2 / νmax)
        κCFL = wizard.Δt / (Δ^2 / κmax)

        @printf("[%06.2f%%] i: %d, t: %.2e, umax: (%.2e, %.2e, %.2e), CFL: %.2e, νκmax: (%.2e, %.2e), νκCFL: (%.2e, %.2e), next Δt: %.2e, wall time: %s\n",
                progress, model.clock.iteration, model.clock.time, umax, vmax, wmax,
                CFL, νmax, κmax, νCFL, κCFL, wizard.Δt, prettytime(simulation.run_time))
    end

    simulation = Simulation(model, Δt=wizard, stop_time=end_time,
                            progress=print_progress, iteration_interval=Ni)
    push!(simulation.output_writers, field_writer, profile_writer, statistics_writer)
    run!(simulation)

    return simulation
end


# sta gives the Spike Triggered Average (average neuron activity preceding a spike)
# stim is the recording of voltage over time in the neuron
# ρ is a matching time series which records the neuron's spikes
# num_timesteps gives the desired time period to record prior to the spike (preset length)
function sta(stim, ρ, num_timesteps)

	# sta_val will store the Spike Triggered Average of the desired preset length
    sta_val = zeros(num_timesteps)

    # Get the indices at the start of each of the presets (ρ records spikes as non-zero values)
    # ignore spikes before num_timesteps as there is not enough data to include those in the average
    start_times = find(ρ[num_timesteps:end])

    # The number of presets
    num_presets = length(start_times)

	# Current positions in the presets (vector)
	current = start_times
	for i = 1:length(sta_val)
		# Average of the current index (there are num_presets)
		sta_val[i] = sum(stim[current]) / num_presets
		current += 1
	end

	return sta_val
end


"""
Sampling from DPP and k-DPP.
References:
-----------
    [1] Kulesza, A., and B. Taskar. Determinantal point processes for machine
        learning. arXiv preprint arXiv:1207.6083, 2012.
    [2] Kang, B. Fast determinantal point process sampling with application to
        clustering. NIPS, 2013.
"""

"""
    _sample_mask!(rng, M, Λ, i)

Sample a mask for an elementary DPP.
"""
function _sample_mask!(
    rng::AbstractRNG, M::AbstractMatrix{Bool}, Λ::AbstractArray{<:Real}, i::Int
)
    for j in 1:length(Λ)
        M[j, i] = rand(rng) < (Λ[j] / (Λ[j] + 1))
    end
    return M
end

"""
    _sample_k_mask!(rng, M, Λ, E, k, i)

Sample a mask for an elementary k-DPP.
"""
function _sample_k_mask!(
    rng::AbstractRNG,
    M::AbstractMatrix{Bool},
    Λ::AbstractArray{<:Real},
    E::AbstractMatrix{<:Real},
    k::Int,
    i::Int,
)
    j = length(Λ)
    remaining = k

    # iteratively sample a k-mask
    while remaining > 0
        # compute marginal of j given that we choose remaining values from 1:j
        if j == remaining
            marg = 1
        else
            marg = Λ[j] * E[remaining, j] / E[remaining + 1, j + 1]
        end

        # sample marginal
        if rand(rng) <= marg
            M[j, i] = true
            remaining -= 1
        end
        j -= 1
    end
    return M
end

"""
    _sample_from_elementary(rng, V, M, i)

Exact sampling from an elementary DPP. The algorithm based on [1].
"""
function _sample_from_elementary(
    rng::AbstractRNG, V::AbstractMatrix{T}, M::AbstractMatrix{Bool}, i::Int
) where {T<:Real}
    # select the elementary DPP
    V_mask = M[:, i]

    # edge case: empty sample
    if isempty(V_mask) # !any(V_mask)
        return Int[]
    end

    # select the kernel of the elementary DPP
    L = V[:, V_mask]

    Y = Int[]
    mask = ones(Bool, size(L, 2))
    prob = zeros(T, size(L, 1))

    for i in 1:size(L, 2)
        # compute probabilities
        fill!(prob, zero(T))
        for c in 1:size(L, 2)
            !mask[c] && continue
            for r in 1:size(L, 1)
                prob[r] += L[r, c] .^ 2
            end
        end
        prob ./= sum(prob)

        # sample a point in the original space
        h = findfirst(rand(rng) .<= cumsum(prob))
        push!(Y, h)

        # select and mask-out an element
        j = get_first_nz_idx(L[h, :], mask)
        mask[j] = false

        if any(mask)
            # Subtract scaled Lj from other columns so that their
            # projections on e_s[i] turns into 0. This operation
            # preserves the rank of L_{-j}.
            for c in 1:size(L, 2)
                !mask[c] && continue
                for r in 1:size(L, 1)
                    L[r, c] -= L[r, j] * L[h, c] / L[h, j]
                end
            end

            # Gram-Schmidt orthogonalization
            L[:, mask] = Matrix(qr(L[:, mask]).Q)
        end
    end

    return sort(Y)
end

Random.rand(pp::PointProcess, n::Int) = rand(GLOBAL_RNG, pp, n)
Random.rand(pp::PointProcess) = rand(GLOBAL_RNG, pp)
Random.rand(rng::AbstractRNG, pp::PointProcess) = first(rand(rng, pp, 1))

"""
    rand([rng::AbstractRNG], dpp::DeterminantalPointProcess)::Vector{Int}
    rand([rng::AbstractRNG], dpp::DeterminantalPointProcess, n::Int)::Vector{Vector{Int}}

Exact sampling from a DPP [1].
Returns a vector of indices with respect to the `L` matrix passed to the `DeterminantalPointProcess`.
The length of each vector can vary from 0 to the `size(L,1)`
"""
function Random.rand(
    rng::AbstractRNG, dpp::DeterminantalPointProcess{T,true}, N::Int) where {T<:Real}
    Λ = SharedArray{T}(dpp.Lfact.values)
    V = SharedMatrix{T}(dpp.Lfact.vectors)
    M = SharedMatrix{Bool}(zeros(Bool, dpp.size, N))

    # step I: sample masks for elementary DPPs
    pmap(
        (i, seed) -> _sample_mask!(MersenneTwister(seed), M, Λ, i),
        1:N,
        abs.(rand(rng, Int, N)),
    )

    # step II: iteratively sample from a mixture of elementary DPPs
    return pmap(
        (i, seed) -> _sample_from_elementary(MersenneTwister(seed), V, M, i),
        1:N,
        abs.(rand(rng, Int, N)),
    )
end

function Random.rand(
    rng::AbstractRNG, dpp::DeterminantalPointProcess{T,false}, N::Int) where {T<:Real}
    Λ = Array{T}(dpp.Lfact.values)
    V = Matrix{T}(dpp.Lfact.vectors)
    M = Matrix{Bool}(zeros(Bool, dpp.size, N))

    # step I: sample masks for elementary DPPs
    map(
        (i, seed) -> _sample_mask!(MersenneTwister(seed), M, Λ, i),
        1:N,
        abs.(rand(rng, Int, N)),
    )

    # step II: iteratively sample from a mixture of elementary DPPs
    return map(
        (i, seed) -> _sample_from_elementary(MersenneTwister(seed), V, M, i),
        1:N,
        abs.(rand(rng, Int, N)),
    )
end

"""
    rand([rng::AbstractRNG], kdpp::kDeterminantalPointProcess)::Vector{Int}
    rand([rng::AbstractRNG], kdpp::kDeterminantalPointProcess, n::Int)::Vector{Vector{Int}}

Exact sampling from a k-DPP [1].
Returns a vector of indices with respect to the `L` matrix passed to the `kDeterminantalPointProcess`
Each vector is of size `k`.
"""
function Random.rand(rng::AbstractRNG, kdpp::kDPP{T,true}, N::Int) where {T<:Real}
    dpp = kdpp.dpp
    Λ = SharedArray{T}(dpp.Lfact.values)
    V = SharedMatrix{T}(dpp.Lfact.vectors)
    M = SharedMatrix{Bool}(zeros(Bool, dpp.size, N))

    # compute elementary symmetric polynomials
    E = SharedMatrix{T}(elem_symm_poly(dpp.Lfact.values, kdpp.k))

    # step I: sample masks for elementary DPPs
    pmap(
        (i, seed) -> _sample_k_mask!(MersenneTwister(seed), M, Λ, E, kdpp.k, i),
        1:N,
        abs.(rand(rng, Int, N)),
    )

    # step II: iteratively sample from a mixture of elementary DPPs
    return pmap(
        (i, seed) -> _sample_from_elementary(MersenneTwister(seed), V, M, i),
        1:N,
        abs.(rand(rng, Int, N)),
    )
end

function Random.rand(rng::AbstractRNG, kdpp::kDPP{T,false}, N::Int) where {T<:Real}
    dpp = kdpp.dpp
    Λ = Array{T}(dpp.Lfact.values)
    V = Matrix{T}(dpp.Lfact.vectors)
    M = Matrix{Bool}(zeros(Bool, dpp.size, N))

    # compute elementary symmetric polynomials
    E = Matrix{T}(elem_symm_poly(dpp.Lfact.values, kdpp.k))

    # step I: sample masks for elementary DPPs
    map(
        (i, seed) -> _sample_k_mask!(MersenneTwister(seed), M, Λ, E, kdpp.k, i),
        1:N,
        abs.(rand(rng, Int, N)),
    )

    # step II: iteratively sample from a mixture of elementary DPPs
    return map(
        (i, seed) -> _sample_from_elementary(MersenneTwister(seed), V, M, i),
        1:N,
        abs.(rand(rng, Int, N)),
    )
end

"""
    Perform one MCMC accept-reject transition for DPP.
"""
function _do_mcmc_step!(rng::AbstractRNG, dpp::DeterminantalPointProcess, state::MCMCState)
    # propose an element to swap
    u = rand(rng, 1:(dpp.size))
    insert = !state[1][u]

    # attempt to make a transition
    if insert
        p = _comp_accept_prob(dpp, state, u, insert)
        rand(rng) < p && _update_mcmc_state!(dpp, state, u, insert)
    else  # delete
        new_state = _update_mcmc_state(dpp, state, u, insert)
        p = _comp_accept_prob(dpp, new_state, u, insert)
        if rand(rng) < p
            state[1] .= new_state[1]
            state[2] .= new_state[2]
        end
    end
end

"""
    Perform one MCMC accept-reject transition for k-DPP.
"""
function _do_mcmc_k_step!(rng::AbstractRNG, kdpp::kDPP, state::MCMCState)
    z, L_z_inv = state

    # propose the elements to swap
    u, v = rand(rng, findall(z)), rand(rng, findall(z .== true))

    # copy the state and delete the u element
    new_state = _update_mcmc_state(kdpp.dpp, state, u, false)

    # attempt to make a transition
    p = _comp_accept_prob(kdpp.dpp, new_state, u, v)
    if rand(rng) < p
        # insert the v element into the new state
        _update_mcmc_state!(kdpp.dpp, new_state, v, true)
        state[1] .= new_state[1]
        state[2] .= new_state[2]
    end
end

"""
    Compute accept probability to insert / delete u from the state.
"""
function _comp_accept_prob(
    dpp::DeterminantalPointProcess, state::MCMCState, u::Int, insert::Bool
)
    z, L_z_inv = state

    d_u = dpp.L[u, u]
    if any(z)
        b_u = dpp.L[z, u]
        d_u -= dot(b_u, L_z_inv[z, z] * b_u)
    end

    return insert ? min(1.0, d_u) : min(1.0, 1.0 / d_u)
end

"""
    Compute accept probability to swap u and v.
"""
function _comp_accept_prob(dpp::DeterminantalPointProcess, state::MCMCState, u::Int, v::Int)
    z, L_z_inv = state

    d_u, d_v = dpp.L[u, u], dpp.L[v, v]
    if any(z)
        b_u, b_v = dpp.L[z, u], dpp.L[z, v]
        d_u -= dot(b_u, L_z_inv[z, z] * b_u)
        d_v -= dot(b_v, L_z_inv[z, z] * b_v)
    end

    return min(1.0, d_v / d_u)
end

"""
    Update the state after u is inserted / deleted.
"""
function _update_mcmc_state!(
    dpp::DeterminantalPointProcess, state::MCMCState, u::Int, insert::Bool
)
    z, L_z_inv = state

    if insert
        d_u = dpp.L[u, u]
        if any(z)
            b_u = dpp.L[z, u]
            x_u = L_z_inv[z, z] * b_u
            d_u -= dot(b_u, x_u)

            L_z_inv[z, z] += (x_u * x_u') / d_u
            L_z_inv[z, u] = L_z_inv[u, z] = -x_u / d_u
        end

        L_z_inv[u, u] = 1.0 / d_u
        z[u] = true
    else  # delete
        z[u] = false

        e = L_z_inv[z, u]
        f = L_z_inv[u, u]

        L_z_inv[z, z] -= (e * e') / f
    end
end

"""
    Update the state after u is inserted / deleted.
"""
function _update_mcmc_state(
    dpp::DeterminantalPointProcess, state::MCMCState, u::Int, insert::Bool
)
    new_state = deepcopy(state)
    _update_mcmc_state!(dpp, new_state, u, insert)
    return new_state
end

"""
    MCMC sampling from a DPP [2].

TODO: Add support for running MCMC in parallel, similar as rand.
          Make sure parallelization produces unbiased and consistent samples.
"""
function randmcmc(
    rng::AbstractRNG,
    dpp::DeterminantalPointProcess{T},
    N::Int;
    init_state=nothing,
    return_final_state::Bool=false,
    mix_eps::Real=1e-1,
    mixing_steps::Int=ceil(Int, dpp.size * log(dpp.size / mix_eps)),
    steps_between_samples::Int=mixing_steps,
) where {T}
    # initialize the Markov chain
    state = init_state
    if state === nothing
        L_z_inv = Array{T}(undef, size(dpp.L))
        z = bitrand(rng, dpp.size)  # TODO: improve initialization (?)
        if any(z)
            L_z_inv[z, z] = pinv(dpp.L[z, z])
        end
        state = (z, L_z_inv)
    end

    # sanity check
    @assert state isa MCMCState

    # mix the Markov chain
    for _ in 1:mixing_steps
        _do_mcmc_step!(rng, dpp, state)
    end

    Y = []
    for _ in 1:N
        push!(Y, findall(state[1]))
        for t in 1:steps_between_samples
            _do_mcmc_step!(rng, dpp, state)
        end
    end

    return return_final_state ? (Y, state) : Y
end

function randmcmc(
    dpp::DeterminantalPointProcess,
    N::Int;
    init_state=nothing,
    return_final_state::Bool=false,
    mix_eps::Real=1e-1,
    mixing_steps::Int=ceil(Int, dpp.size * log(dpp.size / mix_eps)),
    steps_between_samples::Int=mixing_steps,
)
    return randmcmc(
        GLOBAL_RNG,
        dpp,
        N;
        init_state=init_state,
        return_final_state=return_final_state,
        mix_eps=mix_eps,
        mixing_steps=mixing_steps,
        steps_between_samples=steps_between_samples,
    )
end

"""
    randmcmc([rng::AbstractRNG], kdpp::kDPP, N::Int; kwargs...)

MCMC sampling from a k-DPP [2].

## Arguments

- `rng` : Random number generator (by default Random.GLOBAL_RNG is used)
- `kdpp` : k-DeterminantalPointProcess
- `N` : Number of samples

## Keyword Arguments

- `init_state`
- `return_final_state`
- `mix_eps`
- `mixing_steps`
- `steps_between_samples`

TODO:
- Add support for running MCMC in parallel, similar as rand.
- Make sure parallelization produces unbiased and consistent samples.
"""
function randmcmc(
    rng::AbstractRNG,
    kdpp::kDPP{T},
    N::Int;
    init_state=nothing,
    return_final_state::Bool=false,
    mix_eps::Real=1e-1,
    mixing_steps::Int=ceil(Int, kdpp.k * log(kdpp.k / mix_eps)),
    steps_between_samples::Int=mixing_steps,
) where {T}

    # initialize the Markov chain
    state = init_state
    if state === nothing
        L_z_inv = Array{T}(undef, size(kdpp.dpp.L))
        z = falses(kdpp.dpp.size)  # TODO: improve initialization (?)
        z[1:(kdpp.k)] .= true
        if any(z)
            L_z_inv[z, z] = pinv(kdpp.dpp.L[z, z])
        end
        state = (z, L_z_inv)
    end

    # sanity check
    @assert typeof(state) == MCMCState
    @assert sum(state[1]) == kdpp.k

    # mix the Markov chain
    for _ in 1:mixing_steps
        _do_mcmc_k_step!(rng, kdpp, state)
    end

    Y = []
    for _ in 1:N
        push!(Y, findall(state[1]))
        for t in 1:steps_between_samples
            _do_mcmc_k_step!(rng, kdpp, state)
        end
    end

    return return_final_state ? (Y, state) : Y
end

function randmcmc(
    kdpp::kDPP,
    N::Int;
    init_state=nothing,
    return_final_state::Bool=false,
    mix_eps::Real=1e-1,
    mixing_steps::Int=ceil(Int, kdpp.k * log(kdpp.k / mix_eps)),
    steps_between_samples::Int=mixing_steps,
)
    return randmcmc(
        GLOBAL_RNG,
        kdpp,
        N;
        init_state=init_state,
        return_final_state=return_final_state,
        mix_eps=mix_eps,
        mixing_steps=mixing_steps,
        steps_between_samples=steps_between_samples,
    )
end


module SimpleSudoku

export solve_sudoku

# find zeros go through each one and find the ones that are not in the same row, column or box
function get_candidates(grid, row, col)
    grid_id_row = ceil(Int, row / 3)
    grid_id_col = ceil(Int, col / 3)

    candidates = setdiff(0:9, mapreduce(vcat, (
        grid[1:row-1, col],
        grid[row+1:9, col],
        grid[row, 1:col-1],
        grid[row, col+1:9], +
        # the 3x3 grid within which the cell is contained
        grid[(grid_id_row-1)*3+1:grid_id_row*3, (grid_id_col-1)*3+1:grid_id_col*3])) do items
        vec(items)
    end)

    candidates
end

get_candidates(grid, cart) = get_candidates(grid, cart[1], cart[2])

"""
    solve_sudoku(grid)

Args:

    grid a 9x9 integer array where 0 indicates unfilled poistions

Returns:
    Tuple{grid::Array{2}, is_solution_valid::Bool}

    where `grid` is a solved grid if the second return value is true
    If the second return value is false then it means there is
        no solution starting with the input grid
"""
function solve_sudoku(grid)
    @assert size(grid) == (9, 9)
    @assert all(in.(grid, Ref(0:9)))
    # these are the unfilled positions
    zero_pos = findall(==(0), grid)

    if length(zero_pos) == 0
        return grid, true
    end

    cand = get_candidates.(Ref(grid), zero_pos)

    # any of them has no possible candidates; then it's an invalid solution
    if any(length.(cand) .== 0)
        return grid, false
    end

    # find the cell with the lowest number of candidates
    idx = argmin(length.(cand))

    # try the candidates one by one
    zp = zero_pos[idx]
    for c in cand[idx]
        cgrid = copy(grid)
        cgrid[zp] = c
        new_grid, valid = solve_sudoku(cgrid)
        if valid
            return new_grid, true
        end
    end

    return grid, false
end

end


using BenchmarkTools
BenchmarkTools.DEFAULT_PARAMETERS.samples = 100

function compute(n::Int)::Int
    result = 0
    for i ∈ 1:trunc(Int, log10(n))
        for j ∈ 10 ^ trunc(Int, ceil(i / 2 - 1)):10 ^ trunc(Int, ceil(i / 2)) - 1
            palindrome = parse(Int, string(string(j)[1:end - i % 2], reverse(string(j))))
            binary = string(palindrome, base = 2)
            if binary == reverse(binary)
                result += palindrome
            end
        end
    end
    return result
end

compute(1_000_000)

@benchmark compute(1_000_000)


## cholesky.jl : Gaussian random field generator based on the Cholesky decomposition of the covariance matrix

"""
    Cholesky <: GaussianRandomFieldGenerator

A [`GaussiandRandomFieldGenerator`](@ref) based on a Cholesky factorization of the covariance matrix.

# Examples
```jldoctest
julia> m = Matern(0.1,1.0)
Matérn (λ=0.1, ν=1.0, σ=1.0, p=2.0)

julia> c = CovarianceFunction(2,m)
2d Matérn covariance function (λ=0.1, ν=1.0, σ=1.0, p=2.0)

julia> pts1 = 0:0.02:1; pts2 = 0:0.02:1
0.0:0.02:1.0

julia> grf = GaussianRandomField(c,Cholesky(),pts1,pts2)
Gaussian random field with 2d Matérn covariance function (λ=0.1, ν=1.0, σ=1.0, p=2.0) on a 51x51 structured grid, using a Cholesky decomposition

julia> plot(grf)
[...]

```
Note that the Cholesky factorization requires the covariance matrix to be `Symmetric` and positive definite. If the covariance matrix is not `Symmetric`, a warning will be thrown but the method will try to continue with an approximate symmetric matrix. If the covariance matrix is not positive definite, an error will be thrown. Try using the `Spectral` method in that case.

See also: [`Spectral`](@ref), [`KarhunenLoeve`](@ref), [`CirculantEmbedding`](@ref)
"""
struct Cholesky <: GaussianRandomFieldGenerator end

const CholeskyGRF = GaussianRandomField{C,Cholesky} where {C}

# compose a GaussianRandomField using Cholesky factorization
function _GaussianRandomField(mean,cov,method::Cholesky,pts...)
    C = apply(cov,pts,pts)

    # error checking
    issymmetric(C) || warn("to use a Cholesky factorization, the covariance matrix must be symmetric/hermitian")
    isposdef(C) || throw(ArgumentError("to use a Cholesky factorization, the covariance matrix must be positive definite"))

    # compute Cholesky factorization
    L = (cholesky(Symmetric(C))).U'

    GaussianRandomField{typeof(cov),Cholesky,typeof(pts)}(mean,cov,pts,L)
end

# returns the required dimension of the random points
randdim(grf::CholeskyGRF) = size(grf.data,1)

# sample from the GaussianRandomField using Cholesky factorization
function _sample(grf::CholeskyGRF, xi)
	grf.mean + std(grf.cov)*reshape(grf.data*xi,size(grf.mean))
end

show(io::IO,::Cholesky) = print(io,"Cholesky decomposition")


using ChainRulesTestUtils, DifferenceEquations, Distributions, LinearAlgebra, Test, Zygote
using DelimitedFiles
using DiffEqBase
using FiniteDiff: finite_difference_gradient

function joint_likelihood_1(A, B, C, u0, noise, observables, D; kwargs...)
    problem = LinearStateSpaceProblem(A, B, u0, (0, size(observables, 2)); C, observables_noise = D,
                                      noise, observables, kwargs...)
    return solve(problem, DirectIteration()).logpdf
end

# CRTU has problems with generating random MvNormal, so just testing diagonals
function kalman_likelihood(A, B, C, u0, observables, D; kwargs...)
    problem = LinearStateSpaceProblem(A, B, u0, (0, size(observables, 2)); C, observables_noise = D,
                                      u0_prior = MvNormal(u0, diagm(ones(length(u0)))),
                                      noise = nothing, observables, kwargs...)
    return solve(problem).logpdf
end

# Matrices from RBC
A_rbc = [0.9568351489231076 6.209371005755285;
         3.0153731819288737e-18 0.20000000000000007]
B_rbc = reshape([0.0; -0.01], 2, 1) # make sure B is a matrix
C_rbc = [0.09579643002426148 0.6746869652592109; 1.0 0.0]
D_rbc = MvNormal(Diagonal(abs2.([0.1, 0.1])))
u0_rbc = zeros(2)

observables_rbc = readdlm(joinpath(pkgdir(DifferenceEquations), "test/data/RBC_observables.csv"),
                          ',')' |> collect
noise_rbc = readdlm(joinpath(pkgdir(DifferenceEquations), "test/data/RBC_noise.csv"), ',')' |>
            collect
# Data and Noise
T = 5
observables_rbc = observables_rbc[:, 1:T]
noise_rbc = noise_rbc[:, 1:T]

@testset "basic inference" begin
    prob = LinearStateSpaceProblem(A_rbc, B_rbc, u0_rbc, (0, size(observables_rbc, 2)); C = C_rbc,
                                   observables_noise = D_rbc, noise = noise_rbc,
                                   observables = observables_rbc)
    @inferred LinearStateSpaceProblem(A_rbc, B_rbc, u0_rbc, (0, size(observables_rbc, 2));
                                      C = C_rbc, observables_noise = D_rbc, noise = noise_rbc,
                                      observables = observables_rbc)

    sol = solve(prob)
    @inferred solve(prob)

    DiffEqBase.get_concrete_problem(prob, false)
    @inferred DiffEqBase.get_concrete_problem(prob, false)

    joint_likelihood_1(A_rbc, B_rbc, C_rbc, u0_rbc, noise_rbc, observables_rbc, D_rbc)
    @inferred joint_likelihood_1(A_rbc, B_rbc, C_rbc, u0_rbc, noise_rbc, observables_rbc, D_rbc)
end

@testset "basic kalman inference" begin
    prob = LinearStateSpaceProblem(A_rbc, B_rbc, u0_rbc, (0, size(observables_rbc, 2)); C = C_rbc,
                                   observables_noise = D_rbc, observables = observables_rbc,
                                   u0_prior = MvNormal(u0_rbc, diagm(ones(length(u0_rbc)))))
    @inferred LinearStateSpaceProblem(A_rbc, B_rbc, u0_rbc, (0, size(observables_rbc, 2));
                                      C = C_rbc, observables_noise = D_rbc,
                                      observables = observables_rbc,
                                      u0_prior = MvNormal(u0_rbc, diagm(ones(length(u0_rbc)))))

    sol = solve(prob)
    @inferred solve(prob)

    prob_concrete = DiffEqBase.get_concrete_problem(prob, false)
    @inferred DiffEqBase.get_concrete_problem(prob, false)

    kalman_likelihood(A_rbc, B_rbc, C_rbc, u0_rbc, observables_rbc, D_rbc)
    @inferred kalman_likelihood(A_rbc, B_rbc, C_rbc, u0_rbc, observables_rbc, D_rbc)
end

gradient((args...) -> joint_likelihood_1(args..., observables_rbc, D_rbc), A_rbc, B_rbc, C_rbc,
         u0_rbc, noise_rbc)

@testset "linear rbc joint likelihood" begin
    @test joint_likelihood_1(A_rbc, B_rbc, C_rbc, u0_rbc, noise_rbc, observables_rbc, D_rbc) ≈
          -690.9407412360038
    @inferred joint_likelihood_1(A_rbc, B_rbc, C_rbc, u0_rbc, noise_rbc, observables_rbc, D_rbc) #
    gradient((args...) -> joint_likelihood_1(args..., observables_rbc, D_rbc), A_rbc, B_rbc, C_rbc,
             u0_rbc, noise_rbc)

    test_rrule(Zygote.ZygoteRuleConfig(),
               (args...) -> joint_likelihood_1(args..., observables_rbc, D_rbc), A_rbc, B_rbc,
               C_rbc, u0_rbc, noise_rbc; rrule_f = rrule_via_ad, check_inferred = false)
end

gradient((args...) -> kalman_likelihood(args..., observables_rbc, D_rbc), A_rbc, B_rbc, C_rbc,
         u0_rbc)

@testset "linear rbc kalman likelihood" begin
    @test kalman_likelihood(A_rbc, B_rbc, C_rbc, u0_rbc, observables_rbc, D_rbc) ≈
          -607.3698273765538
    @inferred kalman_likelihood(A_rbc, B_rbc, C_rbc, u0_rbc, observables_rbc, D_rbc) # would this catch inference problems in the solve?
    test_rrule(Zygote.ZygoteRuleConfig(),
               (args...) -> kalman_likelihood(args..., observables_rbc, D_rbc), A_rbc, B_rbc, C_rbc,
               u0_rbc; rrule_f = rrule_via_ad, check_inferred = false)
end

# Load FVGQ data for checks
A_FVGQ = readdlm(joinpath(pkgdir(DifferenceEquations), "test/data/FVGQ20_A.csv"), ',')
B_FVGQ = readdlm(joinpath(pkgdir(DifferenceEquations), "test/data/FVGQ20_B.csv"), ',')
C_FVGQ = readdlm(joinpath(pkgdir(DifferenceEquations), "test/data/FVGQ20_C.csv"), ',')
D_FVGQ = MvNormal(Diagonal(abs2.(ones(6) * 1e-3)))

observables_FVGQ = readdlm(joinpath(pkgdir(DifferenceEquations),
                                    "test/data/FVGQ20_observables.csv"), ',')' |> collect

noise_FVGQ = readdlm(joinpath(pkgdir(DifferenceEquations), "test/data/FVGQ20_noise.csv"), ',')' |>
             collect
u0_FVGQ = zeros(size(A_FVGQ, 1))

@testset "linear FVGQ joint likelihood" begin
    @test joint_likelihood_1(A_FVGQ, B_FVGQ, C_FVGQ, u0_FVGQ, noise_FVGQ, observables_FVGQ,
                             D_FVGQ) ≈ -1.4648817357717388e9
    @inferred joint_likelihood_1(A_FVGQ, B_FVGQ, C_FVGQ, u0_FVGQ, noise_FVGQ, observables_FVGQ,
                                 D_FVGQ)
    test_rrule(Zygote.ZygoteRuleConfig(),
               (args...) -> joint_likelihood_1(args..., observables_FVGQ, D_FVGQ), A_FVGQ, B_FVGQ,
               C_FVGQ, u0_FVGQ, noise_FVGQ; rrule_f = rrule_via_ad, check_inferred = false)
end

@testset "linear FVGQ Kalman" begin
    # Note: set rtol to be higher than the default case because of huge gradient numbers
    @test kalman_likelihood(A_FVGQ, B_FVGQ, C_FVGQ, u0_FVGQ, observables_FVGQ, D_FVGQ) ≈
          -108.52706300389917
    gradient((args...) -> kalman_likelihood(args..., observables_FVGQ, D_FVGQ), A_FVGQ, B_FVGQ,
             C_FVGQ, u0_FVGQ)

    # TODO: this is not turned on because the numbers explode.  Need better unit test data to be interior
    # test_rrule(Zygote.ZygoteRuleConfig(), (args...) -> kalman_likelihood(args..., observables, D),
    #            A, B, C, u0; rrule_f = rrule_via_ad, check_inferred = false, rtol = 1e-5)
end


module testMoments
using Test
using LinearAlgebra ; const la = LinearAlgebra

srcDir = dirname(dirname(@__FILE__))*"\\src\\"
include(srcDir*"Moments.jl")
import .Moments ; const mom = Moments

@testset "eᵢ" begin
    n = 7
    for k in 1:n
        e = mom.eᵢ(n,k)
        @test length(e) == n
        @test maximum(e) == 1
        @test minimum(e) == 0
        @test sum(e) == 1
    end
end

@testset "var_kron" begin
    A = repeat([[1,0,1]],2,2)
    B = repeat([[0,1,0]],2,2)
    A_tens_B =  mom.var_kron(A,B)
    @test A_tens_B == repeat([[1,1,1]],4,4)

    n = rand(2:4)
    m = rand(2:5)
    k = n*m
    A = [rand(0:2,5) for i in 1:n, j in 1:n]
    B = [rand(0:2,5) for i in 1:m, j in 1:m]
    A_tens_B =  mom.var_kron(A,B)
    @test size(A_tens_B) == (k,k)

    A =  mom.make_mon_expo(4,(1,1))
    B =  mom.make_mon_expo(4,(1,1);isle = false)
    C = mom.var_kron(A,B)
    for i in 1:4, j in 1:4
        rc(k) = (1+(k-1)*4):(4+(k-1)*4)
        @assert C[rc(i),rc(j)] == [A[i,j]] .+ B
    end

    # Need more
end

@testset "make_mom_expo_mat" begin
    mon_expo_mat =  mom.make_mon_expo(3,(1,1);isle = true)
    @test la.diag(mon_expo_mat) == [ [0, 0, 0],
                                  [2, 0, 0],
                                  [0, 2, 0],
                                  [0, 0, 2]]

    n = 4
    t = 2

    M_vec =  mom.make_mon_expo(n,t)
    M_mat =  mom.make_mon_expo(n,(t,t))

    ## Complex moments

    t = 2
    d = (2,2)
    n = sum(2 .* d)
    M_vec =  mom.make_mon_expo(d,t)
    @assert length(M_vec) == binomial(n+t,t)
    M_mat =  mom.make_mon_expo(d,(t,t))
    @assert M_vec == M_mat[1,:]
end

@testset "make_mom_expo_keys" begin
    n = 3
    t = (1,1)
    mek = mom.make_mon_expo(n,2)
    @test mek == [[0, 0, 0],
                            [1, 0, 0],
                            [0, 1, 0],
                            [0, 0, 1],
                            [2, 0, 0],
                            [1, 1, 0],
                            [1, 0, 1],
                            [0, 2, 0],
                            [0, 1, 1],
                            [0, 0, 2]]

    for i in 1:3
        n = rand(3:6,1)[1]
        t = rand(2:4,1)[1]
        mek = mom.make_mon_expo(n,t*2)
        @test size(unique(mek))[1]  == size(mek)[1]
    end
end

## Real

@testset "make_xxᵀ_tens_yyᵀ" begin
    d = (3,3)
    xxᵀtensyyᵀ = mom.make_xxᵀ_tens_yyᵀ(d)
    @test size(xxᵀtensyyᵀ) == (prod(d),prod(d))

    @test diag(xxᵀtensyyᵀ) ==  [[2, 0, 0, 2, 0, 0],
                                [2, 0, 0, 0, 2, 0],
                                [2, 0, 0, 0, 0, 2],
                                [0, 2, 0, 2, 0, 0],
                                [0, 2, 0, 0, 2, 0],
                                [0, 2, 0, 0, 0, 2],
                                [0, 0, 2, 2, 0, 0],
                                [0, 0, 2, 0, 2, 0],
                                [0, 0, 2, 0, 0, 2]]

    @test sum.(xxᵀtensyyᵀ) == repeat([4],prod(d),prod(d))

    for i in 1:5
        d = rand(2:5,2)
        xxᵀtensyyᵀ = mom.make_xxᵀ_tens_yyᵀ(d)
        halfsum(arr) = [sum(arr[1:d[1]]),sum(arr[d[1]+1:end])]
        @test halfsum.(xxᵀtensyyᵀ) == repeat([[2,2]],prod(d),prod(d))
    end
end

@testset "get_ℝ_block_diagᵀ" begin
    d = (3,3)
    t = (4,4)
    ℝ_block_diag = mom.get_ℝ_block_diag(d,t)

    @test sum([size(ℝ_block_diag[b])[1] for b ∈ keys(ℝ_block_diag)]) == binomial(sum(d)+t[1],t[1])
end

## Complex

@testset "make_xxᵀ⨂yyᵀ" begin
    d = rand(2:5,2)
    xx̄ᵀ_tens_yȳᵀ = mom.make_xxᵀ⨂yyᵀ(d)
end

@test mom.split_expo([1,1,1,2,2,2,3,3,3,4,4,4],3,3) == ([1, 1, 1], [2, 2, 2], [3, 3, 3], [4, 4, 4])

@testset "get_xx̄yȳMM_blocks" begin
    d = (2,2) ; t = (2,2)
    Iᵗ = make_mon_expo(d,t[1])
    Iᵗ_temp = map(x ->  sum.(split_expo(x,d...)),Iᵗ)
    r_list  = map(v -> v[1]-v[2],Iᵗ_temp) ; s_list  = map(v -> v[3]-v[4],Iᵗ_temp)

    Iᵗ_d = Dict()
    for r ∈ -t[1]:t[1], s ∈ -t[1]:t[1]
        T = Iᵗ[(r_list .== r) .& (s_list .== s)]
        isempty(T) ? continue : Iᵗ_d[r,s] = T
    end
    nar = [i  for k in keys(Iᵗ_d) for i in Iᵗ_d[k]]
    @test sort(nar) == sort(Iᵗ)

    d₁ = rand(2:4); t₁ = rand(2:4)
    d = (d₁,d₁) ; t = (t₁,t₁)
    xx̄yȳMM_blocks = mom.get_xx̄yȳMM_blocks(d,t)
    for b in keys(xx̄yȳMM_blocks)
        REF = map(x-> mom.split_expo(x,d...),xx̄yȳMM_blocks[b])
        rp =  sum.(map(v -> v[1]-v[2],REF)) ; sp  = sum.(map(v -> v[3]-v[4],REF))

        @test length(unique(rp)) == 1
        @test length(unique(sp)) == 1
        @test (rp[1,1],sp[1,1]) == (0,0)
    end
end

@testset "get_γδ_dict" begin
    d₁ = rand(2:3); t₁ = rand(2:3)
    d = (d₁,d₁) ; t = (t₁,t₁)
    γδ_dict = Moments.get_γδ_dict(d,t)
    for α ∈ keys(γδ_dict)
        γ,δ = γδ_dict[α]
        @test length(unique(γ+δ)) == 1
        @test unique(γ+δ)[1] == α
    end
end

@testset "get_coef" begin

end

@testset "get_expo" begin
end

@testset "get_ℜℑααᶥββᶥ" begin
end

@testset "get_ℝℂ_coefexpo" begin
end

@testset "get_ℂ_block_diagᵀ" begin

end

end


@testset "Hebisch Problems" begin
    (x, ) = @variables x

    #= ::Package:: =#

    #= ::Title:: =#
    #=Waldek*Hebisch - email*May*2013=#

    #= ::Subsection:: =#
    #=Problem #1=#

    @test_int [(x^6 - x^5 + x^4 - x^3 + 1)*exp(x), x, 25, 871*ℯ^x - 870*ℯ^x*x + 435*ℯ^x*x^2 - 145*ℯ^x*x^3 + 36*ℯ^x*x^4 - 7*ℯ^x*x^5 + ℯ^x*x^6]

    #= ::Subsection:: =#
    #=Problem #2=#

    @test_int [(2 - x^2)*exp(x/(x^2 + 2))/(x^3 + 2*x), x, -5, ExpIntegralEi(x/(2 + x^2))]

    @test_int [(2 + 2*x + 3*x^2 - x^3 + 2*x^4)*exp(x/(2 + x^2))/(x^3 + 2*x), x, -5, ℯ^(x/(2 + x^2))*(2 + x^2) + ExpIntegralEi(x/(2 + x^2))]

    #= ::Subsection:: =#
    #=Problem #3=#

    @test_int [(exp(x) + 1)*(exp(exp(x) + x)/(exp(x) + x)), x, 2, ExpIntegralEi(ℯ^x + x)]

    #= ::Subsection:: =#
    #=Problem #4=#

    @test_int [(x^3 - x^2 - 3*x + 1)*(exp(1/(x^2 - 1))/(x^3 - x^2 - x + 1)), x, -6, ℯ^(1/(-1 + x^2))*(1 + x)]

    #= ::Subsection:: =#
    #=Problem #5=#

    @test_int [(log(x)^2 - 1)*exp(1 + 1/log(x))/log(x)^2, x, 1, x*ℯ^(1 + 1/log(x))]

    @test_int [((x + 1)*log(x)^2  - 1)*exp(x + 1/log(x))/log(x)^2, x, -2, ℯ^(x + 1/log(x))*x]
    end

test/0 Independent test suites/Hebisch Problems.jl