include("euler/euler.jl")
using .Numbers: is_palindrome

function compute(n::Int)::Int
    result = 0
    for i ∈ 10 ^ n ÷ 11 * 11:-11:10 ^ (n - 1)
        for j ∈ 10 ^ n - 1:-1:10 ^ (n - 1)
            candidate = i * j
            if candidate ≤ result
                break
            elseif is_palindrome(candidate)
                result = candidate
            end
        end
    end
    return result
end


#############################################################################
#Copyright (c) 2016 Charles Le Losq
#
#The MIT License (MIT)
#
#Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the #Software without restriction, including without limitation the rights to use, copy, #modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, #and to permit persons to whom the Software is furnished to do so, subject to the #following conditions:
#
#The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
#
#THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, #INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
#
# functions.jl contains several mathematic functions
#
#
#############################################################################

"""
Bootstrap function
    bootstrap(data::Array{Float64}, ese::Array{Float64},nbsample::Int64)

Bootstrap of Raman spectra. We generate new datapoints with the basis of existing data and their standard deviation
"""
function bootstrap(data::Array{Float64}, ese::Array{Float64},nbsample::Int64):
    bootsamples = zeros(size(data)[1],nbsample)
    for i = 1:nbsample
        randn!(bootsamples[:,i]) .* ese[:] + data[:]
    end
    return bootsamples
end

function poly(p::Vector{Float64},x::Array{Float64})
    segments = zeros(size(x)[1],size(p)[1])
    for i = 1:size(p)[1]
        segments[:,i] = p[i].*x[:].^(i-1)
    end
    return sum(segments,2)
end

#For Gaussian peaks in spectra:
function gaussiennes(g_amplitudes::Array{Float64},g_frequency::Array{Float64},g_hwhm::Array{Float64},x::Array{Float64})
    segments = zeros(size(x)[1],size(g_amplitudes)[1])
    for i = 1:size(g_amplitudes)[1]
        segments[:,i] = g_amplitudes[i] .*exp(-log(2) .* ((x[:,1]-g_frequency[i])./g_hwhm[i]).^2)
    end
    return sum(segments,2), segments
end

#The real normal distribution / gaussian function
function normal_dist(nd_amplitudes::Array{Float64},nd_centres::Array{Float64},nd_sigmas::Array{Float64},x::Array{Float64})
    segments = zeros(size(x)[1],size(nd_amplitudes)[1])
    for i = 1:size(nd_amplitudes)[1]
        segments[:,i] = nd_amplitudes[i]./(nd_sigmas[i].*sqrt(2*pi))  .*  exp(- (x[:,1]-nd_centres[i]).^2 ./ (2.*nd_sigmas[i].^2))
    end
    return sum(segments,2), segments
end


using Wavelets, DelimitedFiles, LinearAlgebra, PyPlot, MRIReco

function analyzeImage(x::Vector{T},D::Matrix{T},xsize::NTuple{2,Int64},psize::NTuple{2,Int64};t0::Int64=size(D,2),tol=1e-3) where T
  nx,ny = xsize
  px,py = psize
  x = reshape(x,nx,ny)
  x_pad = repeat(x,2,2)[1:nx+px-1,1:ny+py-1] # pad image using periodic boundary conditions
  α = zeros(T,size(D,2),nx,ny)
  patch = zeros(T,px*py)
  for j=1:ny
    for i=1:nx
      patch[:] .= vec(x_pad[i:i+px-1,j:j+py-1])
      norm(patch)==0 && continue
      # matchingpursuit is contained in Wavelets.jl
      α[:,i,j] .= matchingpursuit(patch, x->D*x, x->transpose(D)*x, tol)
    end
  end
  return vec(α)
end

function synthesizeImage(α::Vector{T},D::Matrix{T},xsize::NTuple{2,Int64},psize::NTuple{2,Int64}) where T
  nx,ny = xsize
  px,py = psize
  x = zeros(T,nx+px-1,ny+py-1)
  α = reshape(α,:,nx,ny)
  for j=1:ny
    for i=1:nx
      x[i:i+px-1,j:j+py-1] .+= reshape(D*α[:,i,j],px,py)
    end
  end
  return vec(x[1:nx,1:ny])/(px*py)
end

function dictOp(D::Matrix{T},xsize::NTuple{2,Int64},psize::NTuple{2,Int64},tol::Float64=1.e-3) where T
  produ = x->analyzeImage(x,D,xsize,psize,tol=tol)
  ctprodu = x->synthesizeImage(x,D,xsize,psize)
  return LinearOperator(prod(xsize)*size(D,2),prod(xsize),false,false
          , produ
          , nothing
          , ctprodu )
end

## To test our method, let us load some simulated data and subsample it

# phantom
img = readdlm("data/mribrain100.tsv")

acqData = AcquisitionData(ISMRMRDFile("data/acqDataBrainSim100.h5"))
nx,ny = acqData.encodingSize

# undersample kspace data
acqData = sample_kspace(acqData, 2.0, "poisson", calsize=25,profiles=false);

## Load the Dictionary, build the sparsifying transform and perform the reconstruction

# load the dictionary
D = ComplexF64.(readdlm("data/brainDict98.tsv"))

# some parameters
px, py = (6,6)  # patch size
K = px*py       # number of atoms

# CS reconstruction using Wavelets
params = Dict{Symbol,Any}()
params[:reco] = "standard"
params[:reconSize] = (nx,ny)
params[:iterations] = 50
params[:λ] = 2.e-2
params[:regularization] = "L1"
params[:sparseTrafo] = dictOp(D,(nx,ny),(px,py),2.e-2)
params[:ρ] = 0.1
params[:solver] = "admm"
params[:absTol] = 1.e-4
params[:relTol] = 1.e-2

img_d = reconstruction(acqData,params)
@info "relative error: $(norm(img-img_d)/norm(img))"

# For comparison, let us perform the same reconstruction as above but with a Wavelet transform

delete!(params, :sparseTrafo)
params[:sparseTrafoName] = "Wavelet"

img_w = reconstruction(acqData,params)
@info "relative error: $(norm(img-img_w)/norm(img))"

# use PyPlot for interactive display
figure(34)
clf()
subplot(2,2,1)
title("Original")
imshow(abs.(img[:,:,1,1,1]))
subplot(2,2,3)
title("Wavelet")
imshow(abs.(img_w[:,:,1,1,1]))
subplot(2,2,4)
title("Dictionary")
imshow(abs.(img_d[:,:,1,1,1]))

# export images
filename = joinpath(dirname(pathof(MRIReco)),"../docs/src/assets/brainOrig.png")
exportImage(filename, abs.(img[:,:,1,1,1]) )
filename = joinpath(dirname(pathof(MRIReco)),"../docs/src/assets/brainWavelet.png")
exportImage(filename, abs.(img_w[:,:,1,1,1]) )
filename = joinpath(dirname(pathof(MRIReco)),"../docs/src/assets/brainDict.png")
exportImage(filename, abs.(img_d[:,:,1,1,1]) )


using Plots
default(legend = false)
x = y = range(-5, 5, length = 40)
zs = zeros(0, 40)
n = 100

@gif for i in range(0, stop = 2π, length = n)
    f(x, y) = sin(x + 10sin(i)) + cos(y)

    # create a plot with 3 subplots and a custom layout
    l = @layout [a{0.7w} b; c{0.2h}]
    p = plot(x, y, f, st = [:surface, :contourf], layout = l)

    # induce a slight oscillating camera angle sweep, in degrees (azimuth, altitude)
    plot!(p[1], camera = (10 * (1 + cos(i)), 40))

    # add a tracking line
    fixed_x = zeros(40)
    z = map(f, fixed_x, y)
    plot!(p[1], fixed_x, y, z, line = (:black, 5, 0.2))
    vline!(p[2], [0], line = (:black, 5))

    # add to and show the tracked values over time
    global zs = vcat(zs, z')
    plot!(p[3], zs, alpha = 0.2, palette = cgrad(:blues).colors)
end


using DiffPrecTest
using Statistics, StatsBase, LinearAlgebra
using SparseArrays
using ProximalBase, CoordinateDescent, CovSel
using Random, Distributions
using JLD

# rep   = parse(Int,ARGS[1])
# ip    = parse(Int,ARGS[2])
# dir = ARGS[3]

alpha = 0.1:0.1:1
pArr = [50, 100, 150]
n = 300
NUM_REP = 1000

bnm = Binomial(1, 0.05)

for ip=1:3
  for ialpha=1:10

    p = pArr[ip]
    Random.seed!(54298)

    # generate X
    Ωx = Matrix{Float64}(I, p, p)
    for i = 1:(p-1)
      for j = (i+1):p
        Ωx[i, j] = rand(bnm) * 0.8
        Ωx[j, i] = Ωx[i, j]
      end
    end
    d = Vector{Float64}(undef, p)
    rand!(Uniform(0.5, 2.5), d)
    d .= sqrt.(d)
    D = Diagonal(d)
    Ωx = D * Ωx * D

    level = maximum( diag(Ωx) ) * sqrt(2. * log(p) / n)

    bc = Bernoulli()
    Ωy = copy(Ωx)
    rd = sample(1:p-1, 8; replace=false)
    sort!(rd)
    Ωy[rd[1], rd[5]] += (2 * rand(bc) - 1) * level
    Ωy[rd[2], rd[6]] += (2 * rand(bc) - 1) * level
    Ωy[rd[3], rd[7]] += (2 * rand(bc) - 1) * level
    Ωy[rd[4], rd[8]] += (2 * rand(bc) - 1) * level

    Ωy[rd[5], rd[1]] = Ωy[rd[1], rd[5]]
    Ωy[rd[6], rd[2]] = Ωy[rd[2], rd[6]]
    Ωy[rd[7], rd[3]] = Ωy[rd[3], rd[7]]
    Ωy[rd[8], rd[4]] = Ωy[rd[4], rd[8]]

    de = max( abs(minimum(eigvals(Ωx))), abs(minimum(eigvals(Ωy))) ) + 0.05
    Ωx = (Ωx + de*I)/(1+de)
    Ωy = (Ωy + de*I)/(1+de)

    Ωy = (1. - alpha[ialpha]) * Ωx + alpha[ialpha] * Ωy
    Δ = Ωx - Ωy

    Σx = inv(Symmetric(Ωx))
    Σy = inv(Symmetric(Ωy))
    dist_X = MvNormal(convert(Matrix, Σx))
    dist_Y = MvNormal(convert(Matrix, Σy))

    rp = div(p * (p+1), 2)
    res = Matrix{DiffPrecResultNormal}(undef, rp, NUM_REP)

    @show p, ialpha
    @time for rep=1:NUM_REP

      # generate data
      Random.seed!(12346 + rep*100 + ialpha)
      X = rand(dist_X, n)'
      Y = rand(dist_Y, n)'

      Sx = Symmetric( X'X / n )
      Sy = Symmetric( Y'Y / n )

      it = 0
      for col=1:p
        for row=col:p
          it = it + 1
          res[it, rep] = DiffPrecTest.estimate(SeparateNormal(), X, Y, row, col; Sx=Sx, Sy=Sy)
        end
      end

    end

    @save "YinXia_power_res_$(ip)_$(ialpha).jld" res
  end
end


using Revise
using ADCME
using NNFEM
using JLD2
using PyCall
using LinearAlgebra
using Random, Distributions

reset_default_graph()

E = 200e3

nntype = "ae_scaled"

#nntype = "piecewise"

idx = 1
H_function = "None"
include("nnutil.jl")

n = 50
X_1d = LinRange(-2.0, 2.0, n)
X = zeros(Float64, n, 1)
for i = 1:n
    X[i, 1] = X_1d[i]
end

ε_μ = 0.0; ε_σ = 0.00
noise = rand(Normal(ε_μ, ε_σ), (n,1))
Y = X * E + noise

loss = constant(0.0)
x = constant(X)
dummy = x

#ε = x; ε0 = x; σ0 = x

y = nn(x, dummy, dummy)

loss += mean((y-Y)^2)

sess = Session(); init(sess)
@show run(sess, loss)
for i = 1:10
    BFGS!(sess, loss, 1000)
    ADCME.save(sess, "Data/NNLearn_$(nntype)_ite$(i).mat")
end


using NonlinearEigenproblems: nep_gallery, compute_Mder
using FEASTSolver

function info(Λ, X, residuals, c, r)
    inside(x) = in_contour(x, c, r)
    in_eig = Λ[inside.(Λ)]
    in_res = residuals[inside.(Λ)]
    print("number inside : ")
    println(size(Λ[inside.(Λ)])[1])
    if sum(inside.(Λ)) > 0
        in_res_conv = in_res[in_res .<= 1e-5]
        in_eig_conv = in_eig[in_res .<= 1e-5]
        print("number inside converged : ")
        println(size(in_eig_conv)[1])
        print("max res inside: ")
        println(maximum(residuals[inside.(Λ)]))
        if size(in_res_conv, 1) > 0
            print("max res inside non spurious: ")
            println(maximum(in_res_conv))
        end
    end
    # display(in_eig_conv)
end

nep = nep_gallery("nlevp_native_loaded_string")
T(x) = compute_Mder(nep, x)

n = size(nep, 1)
C = complex(800.0, 0.0)
R = 790

e, v, res = nlfeast_moments!(T, rand(ComplexF64,n,14), 2^4, 10, c=C, r=R, debug=true, moments=3, ϵ=10e-16, store=true)

info(e,v,res,C,R)


export hist_eig

#Uses the method of Sturm sequences to compute the histogram of eigenvalues of a matrix
#
# Reference
#   Cy P. Chan, "Sturm Sequences and the Eigenvalue Distribution of
#   the Beta-Hermite Random Matrix Ensemble", M.Sc. thesis (MIT), 2007
#
# For better numerical stability, this algorithm is modified so that only the sign
# of the determinant is stored.
#
#For general matrices, this is slower than hist(eigvals(M), bins) and is NOT recommended
function hist_eig{GridPoint <: Number}(M::AbstractMatrix, bins::Vector{GridPoint})
  n = size(M)[1]
  NumBins = length(bins)
  histogram = zeros(NumBins)
  SturmSequence = zeros(n)
  for BinId=1:NumBins
    K = M - bins[BinId]*eye(n)
    #Progression of determinants of lower-right submatrices
    #SturmSequence = [det(K[end-j+1:end,end-j+1:end]) for j=1:n]
    #SturmRatioSequence = [SturmSequence[j+1]/SturmSequence[j] for j=1:n]

    SturmSequenceSigns = [int(sign(det(K[end-j+1:end,end-j+1:end]))) for j=1:n]
    SturmSequenceSigns = [1; SturmSequenceSigns]
    SturmRatioSequence = [SturmSequenceSigns[j+1]/SturmSequenceSigns[j] for j=1:n]
    histogram[BinId] = sum([r < 0 for r in SturmRatioSequence])
  end
  histogram = int(diff([histogram; n]))
end

#Uses the method of Sturm sequences to compute the histogram of eigenvalues of
#a symmetric tridiagonal mmatrix
#
# Reference
#   Cy P. Chan, "Sturm Sequences and the Eigenvalue Distribution of
#   the Beta-Hermite Random Matrix Ensemble", M.Sc. thesis (MIT), 2007
#
function hist_eig{GridPoint <: Number}(M::SymTridiagonal, bins::Vector{GridPoint})
  n = size(M)[1]
  NumBins = length(bins)
  histogram = zeros(NumBins)
  r = zeros(n)
  for BinId=1:NumBins
    a, b = M.dv - bins[BinId], M.ev
    #Formula (2.3)
    r[1] = a[1]
    for i=2:n
      r[i] = a[i] - b[i-1]^2/r[i-1]
    end
    histogram[BinId] = sum([SturmRatio < 0 for SturmRatio in r])
  end
  histogram = int(diff([histogram; n]))
end


#############################################################################
# relative_entropy_epicone.jl
# Constrains τ to
#   τ ⪰ e' * X^{1/2} * logm(X^{1/2}*Y^{-1}*X^{1/2}) * X^{1/2} * e
#
# This function implements the semidefinite programming approximation given in
# the reference below.  Parameters m and k control the accuracy of this
# approximation: m is the number of quadrature nodes to use and k the number
# of square-roots to take. See reference for more details.
#
# All expressions and atoms are subtypes of AbstractExpr.
# Please read expressions.jl first.
#
#REFERENCE
#   Ported from CVXQUAD which is based on the paper: "Semidefinite
#   approximations of matrix logarithm" by Hamza Fawzi, James Saunderson and
#   Pablo A. Parrilo (arXiv:1705.00812)
#############################################################################

struct RelativeEntropyEpiCone
    X::AbstractExpr
    Y::AbstractExpr
    m::Integer
    k::Integer
    e::AbstractMatrix
    size::Tuple{Int,Int}

    function RelativeEntropyEpiCone(
        X::AbstractExpr,
        Y::AbstractExpr,
        m::Integer = 3,
        k::Integer = 3,
        e::AbstractArray = Matrix(1.0 * I, size(X)),
    )
        if size(X) != size(Y)
            throw(DimensionMismatch("X and Y must be the same size"))
        end
        n = size(X)[1]
        if size(X) != (n, n)
            throw(DimensionMismatch("X and Y must be square"))
        end
        if length(size(e)) == 1
            e = reshape(e, (size(e)[1], 1))
        end
        erows, ecols = size(e)
        if ndims(e) != 2 || erows != n
            throw(DimensionMismatch("e matrix must have n rows"))
        end
        return new(X, Y, m, k, e, (ecols, ecols))
    end

    function RelativeEntropyEpiCone(
        X::Value,
        Y::AbstractExpr,
        m::Integer = 3,
        k::Integer = 3,
        e::AbstractArray = Matrix(1.0 * I, size(X)),
    )
        return RelativeEntropyEpiCone(Constant(X), Y, m, k, e)
    end
    function RelativeEntropyEpiCone(
        X::AbstractExpr,
        Y::Value,
        m::Integer = 3,
        k::Integer = 3,
        e::AbstractArray = Matrix(1.0 * I, size(X)),
    )
        return RelativeEntropyEpiCone(X, Constant(Y), m, k, e)
    end
    function RelativeEntropyEpiCone(
        X::Value,
        Y::Value,
        m::Integer = 3,
        k::Integer = 3,
        e::AbstractArray = Matrix(1.0 * I, size(X)),
    )
        return RelativeEntropyEpiCone(Constant(X), Constant(Y), m, k, e)
    end
end

struct RelativeEntropyEpiConeConstraint <: Constraint
    head::Symbol
    id_hash::UInt64
    τ::AbstractExpr
    cone::RelativeEntropyEpiCone

    function RelativeEntropyEpiConeConstraint(
        τ::AbstractExpr,
        cone::RelativeEntropyEpiCone,
    )
        if size(τ) != cone.size
            throw(DimensionMismatch("τ must be size $(cone.size)"))
        end
        id_hash = hash((
            cone.X,
            cone.Y,
            cone.m,
            cone.k,
            cone.e,
            :RelativeEntropyEpiCone,
        ))
        return new(:RelativeEntropyEpiCone, id_hash, τ, cone)
    end

    function RelativeEntropyEpiConeConstraint(
        τ::Value,
        cone::RelativeEntropyEpiCone,
    )
        return RelativeEntropyEpiConeConstraint(Constant(τ), cone)
    end
end

in(τ, cone::RelativeEntropyEpiCone) = RelativeEntropyEpiConeConstraint(τ, cone)

function AbstractTrees.children(constraint::RelativeEntropyEpiConeConstraint)
    return (constraint.τ, constraint.cone.X, constraint.cone.Y)
end

# This negative relative entropy function is matrix convex (arxiv:1705.00812).
# So if X and Y are convex sets, then τ ⪰ -D_op(X || Y) will be a convex set.
function vexity(constraint::RelativeEntropyEpiConeConstraint)
    X = vexity(constraint.cone.X)
    Y = vexity(constraint.cone.Y)
    τ = vexity(constraint.τ)

    # NOTE: can't say X == NotDcp() because the NotDcp constructor prints a warning message.
    if typeof(X) == ConcaveVexity || typeof(X) == NotDcp
        return NotDcp()
    end
    if typeof(Y) == ConcaveVexity || typeof(Y) == NotDcp
        return NotDcp()
    end
    # Copied from vexity(c::GtConstraint)
    vex = ConvexVexity() + (-τ)
    if vex == ConcaveVexity()
        vex = NotDcp()
    end
    return vex
end

function glquad(m)
    # Compute Gauss-Legendre quadrature nodes and weights on [0, 1]
    # Code below is from [Trefethen, "Is Gauss quadrature better than
    # Clenshaw-Curtis?", SIAM Review 2008] and computes the weights and
    # nodes on [-1, 1].
    beta = [0.5 ./ sqrt(1 - (2 * i)^-2) for i in 1:(m-1)] # 3-term recurrence coeffs
    T = diagm(1 => beta, -1 => beta) # Jacobi matrix
    s, V = eigen(T)
    w = 2 * V[1, :] .^ 2 # weights
    # Translate and scale to [0, 1]
    s = (s .+ 1) / 2
    w = w' / 2
    return s, w
end

function conic_form!(
    constraint::RelativeEntropyEpiConeConstraint,
    unique_conic_forms::UniqueConicForms,
)
    if !has_conic_form(unique_conic_forms, constraint)
        X = constraint.cone.X
        Y = constraint.cone.Y
        m = constraint.cone.m
        k = constraint.cone.k
        e = constraint.cone.e
        τ = constraint.τ
        n = size(X)[1]
        r = size(e)[2]

        s, w = glquad(m)

        is_complex =
            sign(X) == ComplexSign() ||
            sign(Y) == ComplexSign() ||
            sign(Constant(e)) == ComplexSign()
        if is_complex
            Z = ComplexVariable(n, n)
            T = [ComplexVariable(r, r) for i in 1:m]
        else
            Z = Variable(n, n)
            T = [Variable(r, r) for i in 1:m]
        end

        conic_form!(
            Z in GeomMeanHypoCone(X, Y, 1 // (2^k), false),
            unique_conic_forms,
        )

        for ii in 1:m
            # Note that we are dividing by w here because it is easier
            # to do this than to do sum w_i T(:,...,:,ii) later (cf. line that
            # involves τ)

            conic_form!(
                [
                    e'*X*e-s[ii]*T[ii]/w[ii] e'*X
                    X*e (1-s[ii])*X+s[ii]*Z
                ] ⪰ 0,
                unique_conic_forms,
            )
        end

        conic_form!((2^k) * sum(T) + τ ⪰ 0, unique_conic_forms)

        cache_conic_form!(
            unique_conic_forms,
            constraint,
            Array{Convex.ConicConstr,1}(),
        )
    end
    return get_conic_form(unique_conic_forms, constraint)
end


module TestDiagrammaticPrograms
using Test

using Catlab, Catlab.Graphs, Catlab.CategoricalAlgebra
using Catlab.Programs.DiagrammaticPrograms
using Catlab.Programs.DiagrammaticPrograms: NamedGraph
using Catlab.Graphs.BasicGraphs: TheoryGraph, TheoryReflexiveGraph
using Catlab.Graphs.BipartiteGraphs: TheoryBipartiteGraph
using Catlab.WiringDiagrams.CPortGraphs: ThCPortGraph

@present TheoryDDS(FreeSchema) begin
  X::Ob
  Φ::Hom(X,X)
end

# Graphs
########

g = @graph begin
  s
  t
  s → t
  s → t
end
@test g == parallel_arrows(NamedGraph{Symbol,Union{Symbol,Nothing}}, 2,
                           V=(vname=[:s,:t],), E=(ename=[nothing,nothing],))

g = @graph NamedGraph{Symbol,Symbol} begin
  x, y
  (f, g): x → y
end
@test g == parallel_arrows(NamedGraph{Symbol,Symbol}, 2,
                           V=(vname=[:x,:y],), E=(ename=[:f,:g],))

tri_parsed = @graph NamedGraph{Symbol,Symbol} begin
  v0, v1, v2
  fst: v0 → v1
  snd: v1 → v2
  comp: v0 → v2
end
tri = @acset NamedGraph{Symbol,Symbol} begin
  V = 3
  E = 3
  src = [1,2,1]
  tgt = [2,3,3]
  vname = [:v0, :v1, :v2]
  ename = [:fst, :snd, :comp]
end
@test tri_parsed == tri

# Categories
############

Δ¹_parsed = @fincat begin
  V, E
  (δ₀, δ₁): V → E
  σ₀: E → V

  σ₀ ∘ δ₀ == id(V)
  σ₀ ∘ δ₁ == id(V)
end
Δ¹_graph = @acset NamedGraph{Symbol,Symbol} begin
  V = 2
  E = 3
  src = [1,1,2]
  tgt = [2,2,1]
  vname = [:V, :E]
  ename = [:δ₀, :δ₁, :σ₀]
end
Δ¹ = FinCat(Δ¹_graph, [ [1,3] => empty(Path, Δ¹_graph, 1),
                        [2,3] => empty(Path, Δ¹_graph, 1) ])
@test Δ¹_parsed == Δ¹

# Functors
##########

# Underlying graph of circular port graph.
F = @finfunctor TheoryGraph ThCPortGraph begin
  V => Box
  E => Wire
  src => src ⨟ box
  tgt => tgt ⨟ box
end
@test F == FinFunctor(Dict(:V => :Box, :E => :Wire),
                      Dict(:src => [:src, :box], :tgt => [:tgt, :box]),
                      TheoryGraph, ThCPortGraph)

# Incomplete definition.
@test_throws ErrorException begin
  @finfunctor TheoryGraph ThCPortGraph begin
    V => Box
    src => src ⨟ box
    tgt => tgt ⨟ box
  end
end

# Failure of functorality.
@test_throws ErrorException begin
  @finfunctor TheoryGraph ThCPortGraph begin
    V => Box
    E => Wire
    src => src
    tgt => tgt
  end
end

# Extra definitions.
@test_throws ErrorException begin
  @finfunctor TheoryGraph TheoryReflexiveGraph begin
    V => Box
    E => Wire
    src => src
    tgt => tgt
    refl => refl
  end
end

# GAT expressions.
F = @finfunctor TheoryDDS TheoryDDS begin
  X => X; Φ => id(X)
end
F′ = @finfunctor TheoryDDS TheoryDDS begin
  X => X; Φ => id{X}
end
@test F == F′

# Diagrams
##########

C = FinCat(TheoryGraph)
F_parsed = @diagram C begin
  v::V
  (e1, e2)::E
  (t: e1 → v)::tgt
  (s: e2 → v)::src
end
J = FinCat(@acset NamedGraph{Symbol,Union{Symbol,Nothing}} begin
  V = 3
  E = 2
  src = [2,3]
  tgt = [1,1]
  vname = [:v, :e1, :e2]
  ename = [:t, :s]
end)
@test dom(F_parsed) == J
F = FinFunctor([:V,:E,:E], [:tgt, :src], J, C)
@test F_parsed == F

F_parsed = @diagram TheoryGraph begin
  v => V
  (e1, e2) => E
  t: e1 → v => tgt
  s: e2 → v => src
end
@test F_parsed == F

F_parsed = @diagram TheoryGraph begin
  v::V
  (e1, e2)::E
  (e1 → v)::tgt
  (e2 → v)::src
end
J_parsed = dom(F_parsed)
@test src(graph(J_parsed)) == src(graph(J))
@test tgt(graph(J_parsed)) == tgt(graph(J))

F_parsed′ = @free_diagram TheoryGraph begin
  v::V
  (e1, e2)::E
  tgt(e1) == v
  v == src(e2)
end
@test F_parsed′ == F_parsed

F = @free_diagram TheoryGraph begin
  (e1, e2)::E
  tgt(e1) == src(e2)
end
@test is_functorial(F)
@test collect_ob(F) == [TheoryGraph[:E], TheoryGraph[:E], TheoryGraph[:V]]
@test collect_hom(F) == [TheoryGraph[:tgt], TheoryGraph[:src]]

F = @diagram TheoryDDS begin
  x::X
  (f: x → x)::(Φ⋅Φ)
end
@test only(collect_ob(F)) == TheoryDDS[:X]
@test only(collect_hom(F)) == compose(TheoryDDS[:Φ], TheoryDDS[:Φ])

# Migrations
############

# Pullback migration
#-------------------

# Graph with reversed edges.
F = @migration TheoryGraph TheoryGraph begin
  V => V
  E => E
  src => tgt
  tgt => src
end
@test F isa DataMigrations.DeltaSchemaMigration
@test F == FinFunctor(Dict(:V => :V, :E => :E),
                      Dict(:src => :tgt, :tgt => :src),
                      TheoryGraph, TheoryGraph)

# Variant where target schema is not given.
F = @migration TheoryGraph begin
  E => E
  V => V
  (src: E → V) => tgt
  (tgt: E → V) => src
end
J = FinCat(parallel_arrows(NamedGraph{Symbol,Union{Symbol,Nothing}}, 2,
                           V=(vname=[:E,:V],), E=(ename=[:src,:tgt],)))
@test F == FinDomFunctor([:E,:V], [:tgt,:src], J, FinCat(TheoryGraph))

# Conjunctive migration
#----------------------

# Graph with edges that are paths of length 2.
F = @migration TheoryGraph TheoryGraph begin
  V => V
  E => @join begin
    v::V
    (e₁, e₂)::E
    (e₁ → v)::tgt
    (e₂ → v)::src
  end
  src => e₁ ⋅ src
  tgt => e₂ ⋅ tgt
end
@test F isa DataMigrations.ConjSchemaMigration
F_E = diagram(ob_map(F, :E))
@test nameof.(collect_ob(F_E)) == [:V, :E, :E]
@test nameof.(collect_hom(F_E)) == [:tgt, :src]
F_tgt = hom_map(F, :tgt)
@test collect_ob(F_tgt) == [(3, TheoryGraph[:tgt])]

# Syntactic variant of above.
F′ = @migration TheoryGraph TheoryGraph begin
  V => V
  E => @join begin
    v::V
    (e₁, e₂)::E
    tgt(e₁) == v
    src(e₂) == v
  end
  src => src(e₁)
  tgt => tgt(e₂)
end
@test F′ == F

# Cartesian product of graph with itself.
F = @migration TheoryGraph TheoryGraph begin
  V => @product (v₁::V; v₂::V)
  E => @product (e₁::E; e₂::E)
  src => (v₁ => e₁⋅src; v₂ => e₂⋅src)
  tgt => (v₁ => e₁⋅tgt; v₂ => e₂⋅tgt)
end
F_V = diagram(ob_map(F, :V))
@test collect_ob(F_V) == fill(TheoryGraph[:V], 2)
F_src = hom_map(F, :src)
@test collect_ob(F_src) == [(1, TheoryGraph[:src]), (2, TheoryGraph[:src])]

# Reflexive graph from graph.
F = @migration TheoryReflexiveGraph TheoryGraph begin
  V => @join begin
    v::V
    ℓ::E
    (s: ℓ → v)::src
    (t: ℓ → v)::tgt
  end
  E => @join begin
    (v₁, v₂)::V
    (ℓ₁, ℓ₂, e)::E
    (s₁: ℓ₁ → v₁)::src
    (t₁: ℓ₁ → v₁)::tgt
    (s₂: ℓ₂ → v₂)::src
    (t₂: ℓ₂ → v₂)::tgt
    (s: e → v₁)::src
    (t: e → v₂)::tgt
  end
  refl => begin
    (v₁, v₂) => v
    (ℓ₁, ℓ₂, e) => ℓ
    (s₁, s₂, s) => s
    (t₁, t₂, t) => t
  end
  src => begin
    v => v₁; ℓ => ℓ₁; s => s₁; t => t₁
  end
  tgt => begin
    v => v₂; ℓ => ℓ₂; s => s₂; t => t₂
  end
end
F_tgt = hom_map(F, :tgt)
@test ob_map(F_tgt, :v) == (2, id(TheoryGraph[:V]))
@test hom_map(F_tgt, :t) |> edges |> only == 4

# Free/initial port graph on a graph.
# This is the left adjoint to the underlying graph functor.
F = @migration TheoryGraph begin
  Box => V
  Wire => E
  InPort => @join begin
    v::V
    e::E
    (t: e → v)::tgt
  end
  OutPort => @join begin
    v::V
    e::E
    (s: e → v)::src
  end
  (in_port_box: InPort → Box) => v
  (out_port_box: OutPort → Box) => v
  (src: Wire → OutPort) => begin
    v => src
  end
  (tgt: Wire → InPort) => begin
    v => tgt
  end
end
F_src = hom_map(F, :src)
@test collect_ob(F_src) == [(1, TheoryGraph[:src]), (1, id(TheoryGraph[:E]))]
@test collect_hom(F_src) == [id(shape(codom(F_src)), 1)]

# Gluing migration
#-----------------

# Coproduct of graph with itself.
F = @migration TheoryGraph TheoryGraph begin
  V => @cases (v₁::V; v₂::V)
  E => @cases (e₁::E; e₂::E)
  src => begin
    e₁ => v₁ ∘ src
    e₂ => v₂ ∘ src
  end
  tgt => begin
    e₁ => v₁ ∘ tgt
    e₂ => v₂ ∘ tgt
  end
end
@test F isa DataMigrations.GlueSchemaMigration
F_V = diagram(ob_map(F, :V))
@test collect_ob(F_V) == fill(TheoryGraph[:V], 2)
F_src = hom_map(F, :src)
@test collect_ob(F_src) == [(1, TheoryGraph[:src]), (2, TheoryGraph[:src])]

# Free reflexive graph on a graph.
F = @migration TheoryReflexiveGraph TheoryGraph begin
  V => V
  E => @cases (e::E; v::V)
  src => (e => src)
  tgt => (e => tgt)
  refl => v
end
F_tgt = hom_map(F, :tgt)
@test collect_ob(F_tgt) == [(1, TheoryGraph[:tgt]), (1, id(TheoryGraph[:V]))]

# Vertices in a graph and their connected components.
F = @migration TheoryGraph begin
  V => V
  Component => @glue begin
    e::E; v::V
    (e → v)::src
    (e → v)::tgt
  end
  (component: V → Component) => v
end
F_C = diagram(ob_map(F, :Component))
@test nameof.(collect_ob(F_C)) == [:E, :V]
@test nameof.(collect_hom(F_C)) == [:src, :tgt]

# Gluc migration
#---------------

# Graph with edges that are paths of length <= 2.
F = @migration TheoryGraph TheoryGraph begin
  V => V
  E => @cases begin
    v => V
    e => E
    path => @join begin
      v::V
      (e₁, e₂)::E
      (e₁ → v)::tgt
      (e₂ → v)::src
    end
  end
  src => begin
    e => src
    path => e₁⋅src
  end
  tgt => begin
    e => tgt
    path => e₂⋅tgt
  end
end
@test ob_map(F, :V) isa DataMigrations.GlucQuery
@test F isa DataMigrations.GlucSchemaMigration
F_src = hom_map(F, :src)
@test collect_ob(shape_map(F_src)) == [1,1,1]
F_src_v, F_src_e, F_src_path = components(diagram_map(F_src))
@test collect_ob(F_src_v) == [(1, id(TheoryGraph[:V]))]
@test collect_ob(F_src_e) == [(1, TheoryGraph[:src])]
@test collect_ob(F_src_path) == [(2, TheoryGraph[:src])]

# Graph with edges that are minimal paths b/w like vertices in bipartite graph.
F = @migration TheoryGraph TheoryBipartiteGraph begin
  V => @cases (v₁::V₁; v₂::V₂)
  E => @cases begin
    e₁ => @join begin
      v₂::V₂; e₁₂::E₁₂; e₂₁::E₂₁
      (e₁₂ → v₂)::tgt₂
      (e₂₁ → v₂)::src₂
    end
    e₂ => @join begin
      v₁::V₁; e₂₁::E₂₁; e₁₂::E₁₂
      (e₂₁ → v₁)::tgt₁
      (e₁₂ → v₁)::src₁
    end
  end
  src => begin
    e₁ => v₁ ∘ (e₁₂ ⋅ src₁)
    e₂ => v₂ ∘ (e₂₁ ⋅ src₂)
  end
  tgt => begin
    e₁ => v₁ ∘ (e₂₁ ⋅ tgt₁)
    e₂ => v₂ ∘ (e₁₂ ⋅ tgt₂)
  end
end
@test ob_map(F, :V) isa DataMigrations.GlucQuery
@test F isa DataMigrations.GlucSchemaMigration
F_src = hom_map(F, :src)
@test collect_ob(shape_map(F_src)) == [1,2]
F_src1, F_src2 = components(diagram_map(F_src))
@test collect_ob(F_src1) == [(2, TheoryBipartiteGraph[:src₁])]
@test collect_ob(F_src2) == [(2, TheoryBipartiteGraph[:src₂])]

# Box product of reflexive graph with itself.
F = @migration TheoryReflexiveGraph TheoryReflexiveGraph begin
  V => @product (v₁::V; v₂::V)
  E => @glue begin
    vv => @product (v₁::V; v₂::V)
    ev => @product (e₁::E; v₂::V)
    ve => @product (v₁::V; e₂::E)
    (refl₁: vv → ev) => (e₁ => v₁⋅refl; v₂ => v₂)
    (refl₂: vv → ve) => (v₁ => v₁; e₂ => v₂⋅refl)
  end
  src => begin
    ev => (v₁ => e₁⋅src; v₂ => v₂)
    ve => (v₁ => v₁; v₂ => e₂⋅src)
  end
  tgt => begin
    ev => (v₁ => e₁⋅tgt; v₂ => v₂)
    ve => (v₁ => v₁; v₂ => e₂⋅tgt)
  end
  refl => vv
end
@test ob_map(F, :V) isa DataMigrations.GlucQuery
@test F isa DataMigrations.GlucSchemaMigration
F_src = hom_map(F, :src)
@test collect_ob(shape_map(F_src)) == [1,1,1]
F_src_vv, F_src_ev, F_src_ve = components(diagram_map(F_src))
@test collect_ob(F_src_ev) == [(1, TheoryReflexiveGraph[:src]),
                               (2, id(TheoryReflexiveGraph[:V]))]

end


function discretize_dyns!(
  F,  # coeff. matrix of the sparse tableau eqs.
  nbset,  # tuple with the following elements.
    # nodes: array of Strings containing independent node names.
    # num_nodes: number of nodes.
    # branches: array of struct Branch.
    # num_branches: sum of branch phases.
    # num_dyns: number of dynamic branches.
  h  # time-step size.
  )  #=

discretizes the time derivative operators appearing in the sparse tableau matrix F given as the first argument. The information of branches, the number of nodes, the number of branches and the time step size are respectively given as the second to fifth argument, and they are used in the discretization.

Author: Taku Noda
Started on Aug. 28, 2020
Update on June 29, 2021

Copyright 2021 Central Research Institute of Electric Power Industry

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

=#

( nodes, num_nodes, branches, num_branches, num_dyns ) = nbset

noff_d = num_nodes + 2*num_branches
noff_v = num_nodes + num_branches
h2 = 0.5*h

for branch in branches
  if branch.type == "L"
    nph = branch.nph
    range = 0:( nph - 1 )
    d_range = ( noff_d + branch.p_de ) .+ range
    v_range = ( noff_v + branch.p_vars ) .+ range
    F[ d_range, v_range ] = sparse( -h2*I, nph, nph )
    F[ d_range, d_range ] = sparse( I, nph, nph )
  elseif branch.type == "C"
    nph = branch.nph
    range = 0:( nph - 1 )
    d_range = ( noff_d + branch.p_de ) .+ range
    i_range = ( num_nodes + branch.p_vars ) .+ range
    F[ d_range, i_range ] = sparse( -h2*I, nph, nph )
    F[ d_range, d_range ] = sparse( I, nph, nph )
  elseif branch.type == "NL"
    dpos = noff_d + branch.p_de
    vpos = noff_v + branch.p_vars
    F[ dpos, vpos ] = -h2
    F[ dpos, dpos ] = 1.0
  elseif branch.type == "NC"
    dpos = noff_d + branch.p_de
    ipos = num_nodes + branch.p_vars
    F[ dpos, vpos ] = -h2
    F[ dpos, dpos ] = 1.0
  end
end

end


using Gadfly
using Main.DistMesh2D

function fh(x, y)::Real
    return 1.0
end

function dc(p)
    return -1 + abs(4 - sqrt(sum(p .^ 2)))
end

function dca(p)
    return sqrt(sum(p .^ 2)) - 2
end

function dcrcc(p)
    return max(drectangle(p, -100.0, 100.0, -100.0, 100.0), - dcircle(p, 20.0, 10.0, 15.0))
end

function eplot(x,y)
    Gadfly.plot(x = x, y = y, Geom.path, Coord.cartesian(fixed = true))
end

#x, y = distmesh2d(dcrcc, fh, [-100.0 -100.0; 100.0 100.0], 7.0, [-100.0 -100.0; -100.0 100.0; 100.0 -100.0; 100.0 100.0])
function fdd(p)
    return dcircle(p, 0.0, 0.0, 1.0)
end
x, y = distmesh2d(fdd, fh, [-1.0 -1.0; 1.0 1.0], 0.2)
eplot(x,y)


using LinearAlgebra
using Random, Distributions, Distances, StatsFuns
using Zygote

include("./kernel.jl")
include("./transform.jl")

struct Model{K<:AbstractKernel, G<:AbstractTransform}
    X::Matrix{Float64}
    Y::Vector{Float64}
    k::K
    g::G
    dims::Int
end

function make_fg(f, unpack)
    return function(F, G, x)
        v, back = pullback(x) do x
            f(unpack(x)...)
        end
        if G !== nothing
            G .= back(1.0)[1]
        end
        if F !== nothing
            return v
        end
        return nothing
    end
end

function logprob(mdl, ℓ, θ, ϵ, λ)
    n = length(mdl.Y)
    L = mdl.X / ℓ
    D = pairwise(SqEuclidean(), L; dims=mdl.dims)
    K = (τ -> radial(mdl.k, θ, τ)).(D) + max(ϵ, 1e-8) * I
    cholK = cholesky(Symmetric(K))
    Z = (y -> transform(mdl.g, λ, y)).(mdl.Y)
    lnJ = logjacobian(mdl.g, λ, mdl.Y)
    lnDet = logdet(cholK)
    return lnJ - lnDet / 2 - dot(Z, cholK \ Z) / 2 - n * log2π / 2
end


struct GeoConst{T,N}
    adv::Advection{T,N}
    #    rho::T
    #    g::T
    #    buoyancy_freq_n::T
    odg_b::T
    #    hv_order::Int
    #    hv_val::T
    coefrsqk::NTuple{N,Array{Complex{T},N}}
    pfftbig::Union{Missing,PrepareFftBig{T}}
    function GeoConst(
        adv::Advection{T,N};
        #        rho::T = 1e3,
        #        g::T = 9.81,
        buoyancy_freq_n::T = 3*2*T(2pi)/(24*3600)*sin(T(pi)/4),
        odg_b::T = T(1e-3),
        #        hv_order::Int = 8,
        #        hv_val::T = T(0),
    ) where {T,N}
        @assert N == 2 "the number of dimension must be 2"
        sz = sizeall(adv)
        coefrsqk = ntuple(x -> zeros(Complex{T}, sz), N)
        pfftbig = if T != Float64
            PrepareFftBig(sizeall(adv), T; numdims = N, dims = ntuple(x -> x, N))
        else
            missing
        end
        veck = vec_k_fft.(adv.t_mesh)
        for ind in CartesianIndices(sz)
            tp = ntuple(i -> veck[i][ind.I[i]], N)
            k = sqrt(sum(tp .^ 2))
            v = k == 0 ? 0 : 1 / k
            for d = 1:N
                coefrsqk[d][ind] = v .* tp[d%2+1]
            end
        end
        coefrsqk[1] .*= im # /buoyancy_freq_n
        coefrsqk[2] .*= -im # /buoyancy_freq_n

        return new{T,N}(
            adv,
            #            rho,
            #            g,
            #            buoyancy_freq_n,
            odg_b,
            #            hv_order,
            #            hv_val,
            coefrsqk,
            pfftbig,
        )
    end
end

struct GeoVar{T,N} <: AbstractExtDataAdv
    gc::GeoConst{T,N}
end

getgeovar(adv::Advection{T,N}; kwargs...) where {T,N} =
    GeoVar{T,N}(GeoConst(adv; kwargs...))

function initdata!(geoc::GeoVar{T,N}, advd::AdvectionData{T,N}) where {T,N}

    mesh_x = advd.adv.t_mesh[1]
    mesh_y = advd.adv.t_mesh[2]
    lx = stop(mesh_x) - start(mesh_x)
    ly = stop(mesh_y) - start(mesh_y)
    x, y = points.(advd.adv.t_mesh)
    ee = 4
    σ = lx / 15 # Length scale close to the Rossby radius

    ## Spatial buoyancy field
    anticyclone(cx, cy) =
        (exp.(-ee .* (x .- cx) .^ 2 ./ 2σ^2) .* transpose(exp.(-(y .- cy) .^ 2 ./ 2σ^2)))

    # Warm anticyclones
    advd.data = anticyclone(lx / 4, ly / 4)
    advd.data .+= anticyclone(3lx / 4, ly / 4)

    # Cold cyclones
    advd.data .-= anticyclone(lx / 4, 3ly / 4)
    advd.data .-= anticyclone(3lx / 4, 3ly / 4)

    # Specify the amplitude of the buoyancy
    advd.data .*= geoc.gc.odg_b

end

function initcoef!(geoc::GeoVar{T,N}, advd::AdvectionData{T,N}) where {T,N}
    pfft = geoc.gc.pfftbig
    dt = getcur_t(advd)
    buf = fftgenall(pfft, advd.data)

    if ismissing(advd.bufcur)
        advd.bufcur = zeros(OpTuple{N,T}, sizeall(advd.adv))
    end

    advd.bufcur .=
        dt * OpTuple.(zip(ntuple(x -> real(ifftgenall(pfft, geoc.gc.coefrsqk[x] .* buf)), N)...))
end
@inline function getalpha(
    ::GeoVar{T,N},
    self::AdvectionData{T},
    indext::CartesianIndex, indbuf::CartesianIndex
) where {T,N,Nsp,Nv}
    st = getst(self)
    ind = CartesianIndex((indbuf.I...,indext.I...)[st.invp])
    return ntuple( x -> self.bufcur[ind][st.invp[x]], st.ndims)
end


# Test whether the numerically computed recursion coefficients are correct
using PolyChaos, Test
import LinearAlgebra: norm
# load config
myfile = open("dataRecCoeffs/config.txt")
ns = parse.(Int,readlines(myfile))

nodes = ns[1]:ns[2]:ns[3]
mus = 0.:0.1:1
albe = 0.:0.2:2
tol = 1e-9

@time for n in nodes
        @testset "Hermite $n nodes" begin
            for m in mus
                myfile = open("dataRecCoeffs/hermite$(n)mu$m.txt")
                αβref = parse.(Float64, readlines(myfile))
                αβcom = rm_hermite(n, m)
                @test isapprox(norm(αβref-[αβcom[1];αβcom[2]],Inf),0.;atol=tol)
                αβcom = coeffs(genHermiteOrthoPoly(n-1, m; addQuadrature=false))
                # αβcom = coeffs(OrthoPoly("genhermite",n-1,Dict(:mu=>m)))
                @test isapprox(norm(αβref-[αβcom[:,1];αβcom[:,2]],Inf),0.;atol=tol)
                close(myfile)
            end
        end
end

@time for n in nodes
        @testset "Logistic $n nodes" begin
            myfile = open("dataRecCoeffs/log$(n).txt")
            αβref = parse.(Float64,readlines(myfile))
            αβcom = rm_logistic(n)
            @test isapprox(norm(αβref-[αβcom[1];αβcom[2]],Inf),0.;atol=tol)
            αβcom = coeffs(LogisticOrthoPoly(n-1, addQuadrature=false))
            @test isapprox(norm(αβref-[αβcom[:,1];αβcom[:,2]],Inf),0.;atol=tol)
            close(myfile)
    end
end

@time for n in nodes
        @testset "Jacobi $n nodes" begin
            for al in albe, be in albe
                    myfile = open("dataRecCoeffs/jac$(n)al$(al)be$(be).txt")
                    αβref = parse.(Float64,readlines(myfile))
                    αβcom = rm_jacobi(n,al,be)
                    @test isapprox(norm(αβref-[αβcom[1];αβcom[2]],Inf),0.;atol=tol)
                    αβcom = coeffs(JacobiOrthoPoly(n-1, al, be, addQuadrature=false))
                    @test isapprox(norm(αβref-[αβcom[:,1];αβcom[:,2]],Inf),0.;atol=tol)
                    @test rm_jacobi(n,al) == rm_jacobi(n,al,al)
                    close(myfile)
            end
    end
end

@time for n in nodes
        @testset "Jacobi01 $n nodes" begin
            for al in albe, be in albe
                    myfile = open("dataRecCoeffs/jac01$(n)al$(al)be$(be).txt")
                    αβref = parse.(Float64,readlines(myfile))
                    αβcom = rm_jacobi01(n,al,be)
                    @test rm_jacobi01(n,al) == rm_jacobi01(n,al,al)
                    @test isapprox(norm(αβref-[αβcom[1];αβcom[2]],Inf),0.;atol=tol)
                    close(myfile)
            end
    end
end

@time for n in nodes
        @testset "Laguerre $n nodes" begin
            for m in mus
            myfile = open("dataRecCoeffs/laguerre$(n)a$m.txt")
            αβref = parse.(Float64,readlines(myfile))
            αβcom = rm_laguerre(n,m)
            @test isapprox(norm(αβref-[αβcom[1];αβcom[2]],Inf),0.;atol=tol)
            αβcom = coeffs(genLaguerreOrthoPoly(n-1, m, addQuadrature=false))
            @test isapprox(norm(αβref-[αβcom[:,1];αβcom[:,2]],Inf),0.;atol=tol)
            close(myfile)
        end
    end
end

meipol = 0.1:0.2:2
@time for n in nodes
        @testset "meixner_pollaczek $n nodes" begin
            for lambda in meipol, phi in meipol
                myfile = open("dataRecCoeffs/meixpol$(n)la$(lambda)phi$(phi).txt")
                αβref = parse.(Float64,readlines(myfile))
                αβcom = rm_meixner_pollaczek(n,lambda,phi)
                @test isapprox(norm(αβref-[αβcom[1];αβcom[2]],Inf),0.;atol=tol)
                αβcom = coeffs(MeixnerPollaczekOrthoPoly(n-1, lambda, phi, addQuadrature=false))
                @test isapprox(norm(αβref-[αβcom[:,1];αβcom[:,2]],Inf),0.;atol=tol)
                close(myfile)
                @test rm_meixner_pollaczek(n,lambda) == rm_meixner_pollaczek(n,lambda,pi/2)
                @test rm_meixner_pollaczek(n,lambda,pi/2)[1] == zeros(n)
            end

    end
end


module SpecialMatricesModule
export computeMatrixKinetic, computeMatrixNuclearAttraction, computeMatrixOverlap
using ..IntegralsModule
using ..BasisModule
using ..Geometry

function computeMatrixKinetic(basis::GaussianBasis)
  return [KineticIntegral(cgb1,cgb2) for cgb1 in basis.contractedBFs, cgb2 in basis.contractedBFs]
end

function computeMatrixOverlap(basis::GaussianBasis)
  return [Overlap(cgb1,cgb2) for cgb1 in basis.contractedBFs, cgb2 in basis.contractedBFs]
end

function computeMatrixNuclearAttraction(basis::GaussianBasis,geo::Geometry)
  return [sum([NuclearAttractionIntegral(cgb1,cgb2,atom) for atom in geo.atoms]) for cgb1 in basis.contractedBFs, cgb2 in basis.contractedBFs]
end

end


# Script used to calculate log-likelihood of multivariate skew t distribution
# for comparison with Julia code. Note that the R script of the same name must be run
# before this, as this script uses its output.

using MultivariateStats, SkewDist

# Load R data
X = readdlm("x.txt", ' ')
Y = readdlm("y.txt", ' ')

n, k = size(Y)
p = size(X,2)

βinit = readdlm("Beta.txt", ' ')
Ωinit = readdlm("Omega.txt", ' ')
αinit = ones(k)
νinit = 4.0

ρ, A, η, logν = SkewDist.dplist2optpar(Ωinit, αinit, νinit)
params = SkewDist.write_params(βinit, ρ, A, η, logν)

println("Parameter vector:", params)
nll, grad = SkewDist.nll_and_grad(params, X, Y)
println("Neg. Log Likelihood: ", nll)
print("Gradient: ", grad, "\n")

βinit = llsq(X,Y; bias=false)
resid = Y - X*βinit


export tone

"""
𝑓 = tone(name; A=440.0)
  convert tone name to frequency, equal temperament
    tone("A4") == A
    tone("C♭0") or tone("Cb0")
    tone("B♯9") or tone("B#9")
"""
function tone(name; A=440.0)
  𝑔=2^(1/12)
  wholetone=Int(name[1])-65   # 'A'=0
  octave=Int(name[end])-48 # '0'=0
  semitone = 2*wholetone
  if wholetone >= 2 #C
    semitone -= 1 # no B♯
    octave -= 1 # octave at C
  end
  if wholetone >= 5 #F
    semitone -= 1 # no E♯
  end
  if length(name) == 3
    if name[2] in ['♭', 'b']
      semitone -= 1
    elseif name[2] in ['♯', '#']
      semitone += 1
    end
  end
  return round(A/16.0 * 𝑔^semitone * 2.0^octave; digits=6)
end


"""
    constructtestfields_2layer(gr)

Constructs flow fields for a 2-layer problem with parameters such that
    q1 = Δψ1 + 25*(ψ2-ψ1),
    q2 = Δψ2 + 25/4*(ψ1-ψ2).
"""
function constructtestfields_2layer(gr)
  x, y = gridpoints(gr)
  k₀, l₀ = 2π/gr.Lx, 2π/gr.Ly # fundamental wavenumbers

  # a set of streafunctions ψ1 and ψ2, ...
  ψ1 = @. 1e-3 * ( 1/4*cos(2k₀*x)*cos(5l₀*y) + 1/3*cos(3k₀*x)*cos(3l₀*y) )
  ψ2 = @. 1e-3 * (     cos(3k₀*x)*cos(4l₀*y) + 1/2*cos(4k₀*x)*cos(2l₀*y) )

  # ... their corresponding PVs q1, q2,
  q1 = @. 1e-3 * ( 1/6 *(   75cos(4k₀*x)*cos(2l₀*y) +   2cos(3k₀*x)*( -43cos(3l₀*y) + 75cos(4l₀*y) ) - 81cos(2k₀*x)*cos(5l₀*y) ) )
  q2 = @. 1e-3 * ( 1/48*( -630cos(4k₀*x)*cos(2l₀*y) + 100cos(3k₀*x)*(    cos(3l₀*y) - 15cos(4l₀*y) ) + 75cos(2k₀*x)*cos(5l₀*y) ) )

  # ... and various derived fields, e.g., ∂ψ1/∂x,
  ψ1x = @. 1e-3 * (     -k₀/2*sin(2k₀*x)*cos(5l₀*y) -       k₀*sin(3k₀*x)*cos(3l₀*y) )
  ψ2x = @. 1e-3 * (      -3k₀*sin(3k₀*x)*cos(4l₀*y) -      2k₀*sin(4k₀*x)*cos(2l₀*y) )
  Δψ2 = @. 1e-3 * ( -25*k₀*l₀*cos(3k₀*x)*cos(4l₀*y) - 10*k₀*l₀*cos(4k₀*x)*cos(2l₀*y) )

  q1x = @. 1e-3 * ( 1/6 *( -4k₀*75sin(4k₀*x)*cos(2l₀*y) - 3k₀*2sin(3k₀*x)*( -43cos(3l₀*y) + 75cos(4l₀*y) ) + 2k₀*81sin(2k₀*x)*cos(5l₀*y) ) )
  q2x = @. 1e-3 * ( 1/48*( 4k₀*630sin(4k₀*x)*cos(2l₀*y) - 3k₀*100sin(3k₀*x)*(    cos(3l₀*y) - 15cos(4l₀*y) ) - 2k₀*75sin(2k₀*x)*cos(5l₀*y) ) )

  Δq1 = @. 1e-3 * (k₀*l₀)*( 1/6 *( -20* 75cos(4k₀*x)*cos(2l₀*y) +   2cos(3k₀*x)*( +18*43cos(3l₀*y) - 25*75cos(4l₀*y) ) +29*81cos(2k₀*x)*cos(5l₀*y) ) )
  Δq2 = @. 1e-3 * (k₀*l₀)*( 1/48*( +20*630cos(4k₀*x)*cos(2l₀*y) + 100cos(3k₀*x)*(    -18cos(3l₀*y) + 25*15cos(4l₀*y) ) -29*75cos(2k₀*x)*cos(5l₀*y) ) )

  return ψ1, ψ2, q1, q2, ψ1x, ψ2x, q1x, q2x, Δψ2, Δq1, Δq2
end

"""
    constructtestfields_3layer(gr)

Constructs flow fields for a 3-layer problem with parameters such that
    q1 = Δψ1 + 20ψ2 - 20ψ1,
    q2 = Δψ2 + 20ψ1 - 44ψ2 + 24ψ3,
    q2 = Δψ2 + 12ψ2 - 12ψ3.
"""
function constructtestfields_3layer(gr)
  x, y = gridpoints(gr)
  k₀, l₀ = 2π/gr.Lx, 2π/gr.Ly # fundamental wavenumbers

  # a set of streafunctions ψ1, ψ2, ψ3, ...
  ψ1 = @. 1e-3 * ( 1/4*cos(2k₀*x)*cos(5l₀*y) + 1/3*cos(3k₀*x)*cos(3l₀*y) )
  ψ2 = @. 1e-3 * (     cos(3k₀*x)*cos(4l₀*y) + 1/2*cos(4k₀*x)*cos(2l₀*y) )
  ψ3 = @. 1e-3 * (     cos(1k₀*x)*cos(3l₀*y) + 1/2*cos(2k₀*x)*cos(2l₀*y) )

  Δψ1 = @. -1e-3 * ( 1/4*((2k₀)^2+(5l₀)^2)*cos(2k₀*x)*cos(5l₀*y) + 1/3*((3k₀)^2+(3l₀)^2)*cos(3k₀*x)*cos(3l₀*y) )
  Δψ2 = @. -1e-3 * (     ((3k₀)^2+(4l₀)^2)*cos(3k₀*x)*cos(4l₀*y) + 1/2*((4k₀)^2+(2l₀)^2)*cos(4k₀*x)*cos(2l₀*y) )
  Δψ3 = @. -1e-3 * (     ((1k₀)^2+(3l₀)^2)*cos(1k₀*x)*cos(3l₀*y) + 1/2*((2k₀)^2+(2l₀)^2)*cos(2k₀*x)*cos(2l₀*y) )

  # ... their corresponding PVs q1, q2, q3,
  q1 = @. Δψ1 + 20ψ2 - 20ψ1
  q2 = @. Δψ2 + 20ψ1 - 44ψ2 + 24ψ3
  q3 = @. Δψ3        + 12ψ2 - 12ψ3

  return ψ1, ψ2, ψ3, q1, q2, q3
end

"""
    test_pvtofromstreamfunction_2layer(dev)

Tests the pvfromstreamfunction function that gives qh from ψh. To do so, it
constructs a 2-layer problem with parameters such that
    q1 = Δψ1 + 25*(ψ2-ψ1), q2 = Δψ2 + 25/4*(ψ1-ψ2).
Then given ψ1 and ψ2 it tests if pvfromstreamfunction gives the expected
q1 and q2. Similarly, that streamfunctionfrompv gives ψ1 and ψ2 from q1 and q2.
"""
function test_pvtofromstreamfunction_2layer(dev::Device=CPU())
  n, L = 128, 2π
  gr = TwoDGrid(dev, n, L)

  nlayers = 2      # these choice of parameters give the
  f₀, g = 1, 1     # desired PV-streamfunction relations
  H = [0.2, 0.8]   # q1 = Δψ1 + 25*(ψ2-ψ1), and
  ρ = [4.0, 5.0]   # q2 = Δψ2 + 25/4*(ψ1-ψ2).

  prob = MultiLayerQG.Problem(nlayers, dev; nx=n, Lx=L, f₀=f₀, g=g, H=H, ρ=ρ)
  sol, cl, pr, vs, gr = prob.sol, prob.clock, prob.params, prob.vars, prob.grid

  ψ1, ψ2, q1, q2, ψ1x, ψ2x, q1x, q2x, Δψ2, Δq1, Δq2 = constructtestfields_2layer(gr)

  vs.ψh[:, :, 1] .= rfft(ψ1)
  vs.ψh[:, :, 2] .= rfft(ψ2)

  MultiLayerQG.pvfromstreamfunction!(vs.qh, vs.ψh, pr, gr)
  MultiLayerQG.invtransform!(vs.q, vs.qh, pr)

  vs.qh[:, :, 1] .= rfft(q1)
  vs.qh[:, :, 2] .= rfft(q2)

  MultiLayerQG.streamfunctionfrompv!(vs.ψh, vs.qh, pr, gr)
  MultiLayerQG.invtransform!(vs.ψ, vs.ψh, pr)

  return isapprox(q1, vs.q[:, :, 1], rtol=rtol_multilayerqg) && isapprox(q2, vs.q[:, :, 2], rtol=rtol_multilayerqg) && isapprox(ψ1, vs.ψ[:, :, 1], rtol=rtol_multilayerqg) && isapprox(ψ2, vs.ψ[:, :, 2], rtol=rtol_multilayerqg)
end

"""
    test_pvtofromstreamfunction_3layer()

Tests the pvfromstreamfunction function that gives qh from ψh. To do so, it
constructs a 3-layer problem with parameters such that
  q1 = Δψ1 + 20ψ2 - 20ψ1, q2 = Δψ2 + 20ψ1 - 44ψ2 + 24ψ3, q3 = Δψ3 + 12ψ2 - 12ψ3.
Then given ψ1, ψ2, and ψ2 it tests if pvfromstreamfunction gives the expected
q1, q2, and q3. Similarly, that streamfunctionfrompv gives ψ1, ψ2, and ψ2 from
q1, q2, and q3.
"""
function test_pvtofromstreamfunction_3layer(dev::Device=CPU())
  n, L = 128, 2π
  gr = TwoDGrid(dev, n, L)

  nlayers = 3            # these choice of parameters give the
  f₀, g = 1, 1           # desired PV-streamfunction relations
  H = [0.25, 0.25, 0.5]  # q1 = Δψ1 + 20ψ2 - 20ψ1,
  ρ = [4.0, 5.0, 6.0]    # q2 = Δψ2 + 20ψ1 - 44ψ2 + 24ψ3,
                         # q3 = Δψ3        + 12ψ2 - 12ψ3.

  prob = MultiLayerQG.Problem(nlayers, dev; nx=n, Lx=L, f₀=f₀, g=g, H=H, ρ=ρ)
  sol, cl, pr, vs, gr = prob.sol, prob.clock, prob.params, prob.vars, prob.grid

  ψ1, ψ2, ψ3, q1, q2, q3 = constructtestfields_3layer(gr)

  vs.ψh[:, :, 1] .= rfft(ψ1)
  vs.ψh[:, :, 2] .= rfft(ψ2)
  vs.ψh[:, :, 3] .= rfft(ψ3)

  MultiLayerQG.pvfromstreamfunction!(vs.qh, vs.ψh, pr, gr)
  MultiLayerQG.invtransform!(vs.q, vs.qh, pr)

  vs.qh[:, :, 1] .= rfft(q1)
  vs.qh[:, :, 2] .= rfft(q2)
  vs.qh[:, :, 3] .= rfft(q3)

  MultiLayerQG.streamfunctionfrompv!(vs.ψh, vs.qh, pr, gr)
  MultiLayerQG.invtransform!(vs.ψ, vs.ψh, pr)

  return isapprox(q1, vs.q[:, :, 1], rtol=rtol_multilayerqg) && isapprox(q2, vs.q[:, :, 2], rtol=rtol_multilayerqg) && isapprox(q3, vs.q[:, :, 3], rtol=rtol_multilayerqg) && isapprox(ψ1, vs.ψ[:, :, 1], rtol=rtol_multilayerqg) && isapprox(ψ2, vs.ψ[:, :, 2], rtol=rtol_multilayerqg) && isapprox(ψ3, vs.ψ[:, :, 3], rtol=rtol_multilayerqg)
end

"""
    test_mqg_nonlinearadvection(dt, stepper, dev; kwargs...)

Tests the advection term by timestepping a test problem with timestep dt and
timestepper identified by the string stepper. The test 2-layer problem is
derived by picking a solution q1f, q2f (with streamfunctions ψ1f, ψ2f) for which
the advection terms J(ψn, qn) are non-zero. Next, a forcing Ff is derived such
that a solution to the problem forced by this Ff is then qf.
(This solution may not be realized, at least at long times, if it is unstable.)
"""
function test_mqg_nonlinearadvection(dt, stepper, dev::Device=CPU(); n=128, L=2π, nlayers=2, μ=0.0, ν=0.0, nν=1)

  A = ArrayType(dev)

  tf = 0.5
  nt = round(Int, tf/dt)

  nx, ny = 64, 66
  Lx, Ly = 2π, 2π
  gr = TwoDGrid(dev, nx, Lx, ny, Ly)

  x, y = gridpoints(gr)
  k₀, l₀ = 2π/gr.Lx, 2π/gr.Ly # fundamental wavenumbers

  nlayers = 2       # these choice of parameters give the
  f₀, g = 1, 1      # desired PV-streamfunction relations
  H = [0.2, 0.8]    # q1 = Δψ1 + 25*(ψ2-ψ1), and
  ρ = [4.0, 5.0]    # q2 = Δψ2 + 25/4*(ψ1-ψ2).

  β = 0.35

  U1, U2 = 0.1, 0.05
  u1 = @. 0.5sech(gr.y/(Ly/15))^2; u1 = A(reshape(u1, (1, gr.ny)))
  u2 = @. 0.02cos(3l₀*gr.y);       u2 = A(reshape(u2, (1, gr.ny)))
  uyy1 = real.(ifft( -gr.l.^2 .* fft(u1) ))
  uyy2 = real.(ifft( -gr.l.^2 .* fft(u2) ))

  U = zeros(ny, nlayers)
  CUDA.@allowscalar U[:, 1] = u1 .+ U1
  CUDA.@allowscalar U[:, 2] = u2 .+ U2

  μ, ν, nν = 0.1, 0.05, 1

  η0, σx, σy = 1.0, Lx/25, Ly/20
   η = @. η0*exp( -(x+Lx/8)^2/(2σx^2) - (y-Ly/8)^2/(2σy^2) )
  ηx = @. -(x + Lx/8)/(σx^2) * η

  ψ1, ψ2, q1, q2, ψ1x, ψ2x, q1x, q2x, Δψ2, Δq1, Δq2 = constructtestfields_2layer(gr)

  Ff1 = FourierFlows.jacobian(ψ1, q1, gr)     + @. (β - uyy1 -   25*(U2+u2-U1-u1) )*ψ1x + (U1+u1)*q1x - ν*Δq1
  Ff2 = FourierFlows.jacobian(ψ2, q2 + η, gr) + @. (β - uyy2 - 25/4*(U1+u1-U2-u2) )*ψ2x + (U2+u2)*(q2x + ηx) + μ*Δψ2 - ν*Δq2

  T = eltype(gr)

  Ff = zeros(dev, T, (gr.nx, gr.ny, nlayers))
  @views Ff[:, :, 1] = Ff1
  @views Ff[:, :, 2] = Ff2

  Ffh = zeros(dev, Complex{T}, (gr.nkr, gr.nl, nlayers))
  @views Ffh[:, :, 1] = rfft(Ff1)
  @views Ffh[:, :, 2] = rfft(Ff2)

  function calcFq!(Fqh, sol, t, cl, v, p, g)
    Fqh .= Ffh
    nothing
  end

  prob = MultiLayerQG.Problem(nlayers, dev; nx=nx, ny=ny, Lx=Lx, Ly=Ly,
   f₀=f₀, g=g, H=H, ρ=ρ, U=U, eta=η, β=β, μ=μ, ν=ν, nν=nν, calcFq=calcFq!, stepper=stepper, dt=dt)
  sol, cl, pr, vs, gr = prob.sol, prob.clock, prob.params, prob.vars, prob.grid

  qf = zeros(dev, T, (gr.nx, gr.ny, nlayers))
  @views qf[:, :, 1] = q1
  @views qf[:, :, 2] = q2

  ψf = zeros(dev, T, (gr.nx, gr.ny, nlayers))
  @views ψf[:, :, 1] = ψ1
  @views ψf[:, :, 2] = ψ2

  MultiLayerQG.set_q!(prob, qf)

  stepforward!(prob, nt)
  MultiLayerQG.updatevars!(prob)

  return isapprox(vs.q, qf, rtol=rtol_multilayerqg) && isapprox(vs.ψ, ψf, rtol=rtol_multilayerqg)
end

"""
    test_mqg_linearadvection(dt, stepper, dev; kwargs...)

Tests the advection term of the linearized equations by timestepping a test
problem with timestep dt and timestepper identified by the string stepper.
The test 2-layer problem is derived by picking a solution q1f, q2f (with
streamfunctions ψ1f, ψ2f) for which the advection terms J(ψn, qn) are non-zero.
Next, a forcing Ff is derived such that a solution to the problem forced by this
Ff is then qf. (This solution may not be realized, at least at long times, if it
is unstable.)
"""
function test_mqg_linearadvection(dt, stepper, dev::Device=CPU(); n=128, L=2π, nlayers=2, μ=0.0, ν=0.0, nν=1)

  A = ArrayType(dev)

  tf = 0.5
  nt = round(Int, tf/dt)

  nx, ny = 64, 66
  Lx, Ly = 2π, 2π
  gr = TwoDGrid(dev, nx, Lx, ny, Ly)

  x, y = gridpoints(gr)
  k₀, l₀ = 2π/gr.Lx, 2π/gr.Ly # fundamental wavenumbers

  nlayers = 2       # these choice of parameters give the
  f₀, g = 1, 1      # desired PV-streamfunction relations
  H = [0.2, 0.8]    # q1 = Δψ1 + 25*(ψ2-ψ1), and
  ρ = [4.0, 5.0]    # q2 = Δψ2 + 25/4*(ψ1-ψ2).

  β = 0.35

  U1, U2 = 0.1, 0.05
  u1 = @. 0.5sech(gr.y/(Ly/15))^2; u1 = A(reshape(u1, (1, gr.ny)))
  u2 = @. 0.02cos(3l₀*gr.y);       u2 = A(reshape(u2, (1, gr.ny)))
  uyy1 = real.(ifft( -gr.l.^2 .* fft(u1) ))
  uyy2 = real.(ifft( -gr.l.^2 .* fft(u2) ))

  U = zeros(ny, nlayers)
  CUDA.@allowscalar U[:, 1] = u1 .+ U1
  CUDA.@allowscalar U[:, 2] = u2 .+ U2

  μ, ν, nν = 0.1, 0.05, 1

  η0, σx, σy = 1.0, Lx/25, Ly/20
  η = @. η0*exp( -(x + Lx/8)^2/(2σx^2) -(y-Ly/8)^2/(2σy^2) )
  ηx = @. -(x + Lx/8)/(σx^2) * η

  ψ1, ψ2, q1, q2, ψ1x, ψ2x, q1x, q2x, Δψ2, Δq1, Δq2 = constructtestfields_2layer(gr)

  Ff1 = (β .- uyy1 .- 25*(U2.+u2.-U1.-u1) ).*ψ1x + (U1.+u1).*q1x - ν*Δq1
  Ff2 = FourierFlows.jacobian(ψ2, η, gr) + (β .- uyy2 .- 25/4*(U1.+u1.-U2.-u2) ).*ψ2x + (U2.+u2).*(q2x + ηx) + μ*Δψ2 - ν*Δq2

  T = eltype(gr)

  Ff = zeros(dev, T, (gr.nx, gr.ny, nlayers))
  @views Ff[:, :, 1] = Ff1
  @views Ff[:, :, 2] = Ff2

  Ffh = zeros(dev, Complex{T}, (gr.nkr, gr.nl, nlayers))
  @views Ffh[:, :, 1] = rfft(Ff1)
  @views Ffh[:, :, 2] = rfft(Ff2)

  function calcFq!(Fqh, sol, t, cl, v, p, g)
    Fqh .= Ffh
    nothing
  end

  prob = MultiLayerQG.Problem(nlayers, dev; nx=nx, ny=ny, Lx=Lx, Ly=Ly, f₀=f₀, g=g, H=H, ρ=ρ, U=U, eta=η, β=β, μ=μ, ν=ν, nν=nν, calcFq=calcFq!, stepper=stepper, dt=dt, linear=true)
  sol, cl, pr, vs, gr = prob.sol, prob.clock, prob.params, prob.vars, prob.grid

  qf = zeros(dev, T, (gr.nx, gr.ny, nlayers))
  @views qf[:, :, 1] = q1
  @views qf[:, :, 2] = q2

  ψf = zeros(dev, T, (gr.nx, gr.ny, nlayers))
  @views ψf[:, :, 1] = ψ1
  @views ψf[:, :, 2] = ψ2

  MultiLayerQG.set_q!(prob, qf)

  stepforward!(prob, nt)
  MultiLayerQG.updatevars!(prob)

  return isapprox(vs.q, qf, rtol=rtol_multilayerqg) && isapprox(vs.ψ, ψf, rtol=rtol_multilayerqg)
end

"""
    test_mqg_energies(dev; kwargs...)

Tests the kinetic (KE) and potential (PE) energies function by constructing a
2-layer problem and initializing it with a flow field whose KE and PE are known.
"""
function test_mqg_energies(dev::Device=CPU(); dt=0.001, stepper="ForwardEuler", n=128, L=2π, nlayers=2, μ=0.0, ν=0.0, nν=1)
  nx, ny = 64, 66
  Lx, Ly = 2π, 2π
  gr = TwoDGrid(dev, nx, Lx, ny, Ly)

  x, y = gridpoints(gr)
  k₀, l₀ = 2π/gr.Lx, 2π/gr.Ly # fundamental wavenumbers

  nlayers = 2       # these choice of parameters give the
  f₀, g = 1, 1      # desired PV-streamfunction relations
  H = [0.2, 0.8]    # q1 = Δψ1 + 25*(ψ2-ψ1), and
  ρ = [4.0, 5.0]    # q2 = Δψ2 + 25/4*(ψ1-ψ2).

  prob = MultiLayerQG.Problem(nlayers, dev; nx=nx, ny=ny, Lx=Lx, Ly=Ly, f₀=f₀, g=g, H=H, ρ=ρ)
  sol, cl, pr, vs, gr = prob.sol, prob.clock, prob.params, prob.vars, prob.grid

  ψ1, ψ2, q1, q2, ψ1x, ψ2x, q1x, q2x, Δψ2, Δq1, Δq2 = constructtestfields_2layer(gr)

  qf = zeros(gr.nx, gr.ny, nlayers)
  CUDA.@allowscalar @views qf[:, :, 1] .= q1
  CUDA.@allowscalar @views qf[:, :, 2] .= q2

  MultiLayerQG.set_q!(prob, qf)

  KE, PE = MultiLayerQG.energies(prob)

  return isapprox(KE[1], 61/640*1e-6, rtol=rtol_multilayerqg) && isapprox(KE[2], 3*1e-6, rtol=rtol_multilayerqg) && isapprox(PE[1], 1025/1152*1e-6, rtol=rtol_multilayerqg) && MultiLayerQG.addforcing!(prob.timestepper.RHS₁, sol, cl.t, cl, vs, pr, gr)==nothing
end

function test_mqg_energysinglelayer(dev::Device=CPU(); dt=0.001, stepper="ForwardEuler", nlayers=1, μ=0.0, ν=0.0, nν=1)
  nx, Lx  = 64, 2π
  ny, Ly  = 64, 3π
  gr = TwoDGrid(dev, nx, Lx, ny, Ly)

  x, y = gridpoints(gr)
  k₀, l₀ = 2π/gr.Lx, 2π/gr.Ly # fundamental wavenumbers

  energy_calc = 29/9

  ψ0 = @. sin(2k₀*x)*cos(2l₀*y) + 2sin(k₀*x)*cos(3l₀*y)
  q0 = @. -((2k₀)^2+(2l₀)^2)*sin(2k₀*x)*cos(2l₀*y) - (k₀^2+(3l₀)^2)*2sin(k₀*x)*cos(3l₀*y)

  prob = MultiLayerQG.Problem(nlayers, dev; nx=nx, Lx=Lx, ny=ny, Ly=Ly, stepper=stepper, U=zeros(ny))

  MultiLayerQG.set_q!(prob, reshape(q0, (nx, ny, nlayers)))

  energyq0 = MultiLayerQG.energies(prob)

  return isapprox(energyq0, energy_calc, rtol=rtol_multilayerqg)
end

"""
    test_mqg_fluxes(dt, stepper; kwargs...)

Tests the lateral and vertical eddy fluxes by constructing a 2-layer problem and
initializing it with a flow field whose fluxes are known.
"""
function test_mqg_fluxes(dev::Device=CPU(); dt=0.001, stepper="ForwardEuler", n=128, L=2π, nlayers=2, μ=0.0, ν=0.0, nν=1)
  nx, ny = 128, 126
  Lx, Ly = 2π, 2π
  gr = TwoDGrid(dev, nx, Lx, ny, Ly)

  x, y = gridpoints(gr)
  k₀, l₀ = 2π/gr.Lx, 2π/gr.Ly # fundamental wavenumbers

  nlayers = 2       # these choice of parameters give the
  f₀, g = 1, 1      # desired PV-streamfunction relations
  H = [0.2, 0.8]    # q1 = Δψ1 + 25*(ψ2-ψ1), and
  ρ = [4.0, 5.0]    # q2 = Δψ2 + 25/4*(ψ1-ψ2).
  U = zeros(ny, nlayers)
  U[:, 1] = @. sech(gr.y / 0.2)^2

  prob = MultiLayerQG.Problem(nlayers, dev; nx=nx, ny=ny, Lx=Lx, Ly=Ly, f₀=f₀, g=g, H=H, ρ=ρ, U=U)
  sol, cl, pr, vs, gr = prob.sol, prob.clock, prob.params, prob.vars, prob.grid

  ψ1 = @. cos(k₀*x) * cos(l₀*y) + sin(k₀*x)
  ψ2 = @. cos(k₀*x + π/10) * cos(l₀*y)
  ψ = zeros(gr.nx, gr.ny, nlayers)
  CUDA.@allowscalar @views ψ[:, :, 1] .= ψ1
  CUDA.@allowscalar @views ψ[:, :, 2] .= ψ2
  MultiLayerQG.set_ψ!(prob, ψ)
  lateralfluxes, verticalfluxes = MultiLayerQG.fluxes(prob)

  return CUDA.@allowscalar isapprox(lateralfluxes[1], 0.00626267, rtol=1e-6) && CUDA.@allowscalar isapprox(lateralfluxes[2], 0, atol=1e-12) && CUDA.@allowscalar isapprox(verticalfluxes[1], -0.196539, rtol=1e-6)
end

"""
    test_mqg_fluxessinglelayer(dt, stepper; kwargs...)

Tests the lateral eddy fluxes by constructing a 1-layer problem and initializing
it with a flow field whose fluxes are known.
"""
function test_mqg_fluxessinglelayer(dev::Device=CPU(); dt=0.001, stepper="ForwardEuler", n=128, L=2π, μ=0.0, ν=0.0, nν=1)
  nlayers = 1

  nx, ny = 128, 126
  Lx, Ly = 2π, 2π
  gr = TwoDGrid(dev, nx, Lx, ny, Ly)

  x, y = gridpoints(gr)
  k₀, l₀ = 2π/gr.Lx, 2π/gr.Ly # fundamental wavenumbers

  U = zeros(ny, nlayers)
  U = @. sech(gr.y / 0.2)^2

  prob = MultiLayerQG.Problem(nlayers, dev; nx=nx, ny=ny, Lx=Lx, Ly=Ly, U=U)
  sol, cl, pr, vs, gr = prob.sol, prob.clock, prob.params, prob.vars, prob.grid

  ψ = @. cos(k₀*x) * cos(l₀*y) + sin(k₀*x)
  MultiLayerQG.set_ψ!(prob, ψ)
  lateralfluxes = MultiLayerQG.fluxes(prob)

  return CUDA.@allowscalar isapprox(lateralfluxes[1], 0.0313134, atol=1e-7)
end

"""
    test_setqsetψ(dt, stepper; kwargs...)

Tests the set_q!() and set_ψ!() functions that initialize sol with a flow with
given `q` or `ψ` respectively.
"""
function test_mqg_setqsetψ(dev::Device=CPU(); dt=0.001, stepper="ForwardEuler", n=64, L=2π, nlayers=2, μ=0.0, ν=0.0, nν=1)
  nx, ny = 32, 34
  L = 2π
  gr = TwoDGrid(dev, nx, L, ny, L)

  x, y = gridpoints(gr)
  k₀, l₀ = 2π/gr.Lx, 2π/gr.Ly # fundamental wavenumbers

  nlayers = 2       # these choice of parameters give the
  f₀, g = 1, 1      # desired PV-streamfunction relations
  H = [0.2, 0.8]    # q1 = Δψ1 + 25*(ψ2-ψ1), and
  ρ = [4.0, 5.0]    # q2 = Δψ2 + 25/4*(ψ1-ψ2).

  prob = MultiLayerQG.Problem(nlayers, dev; nx=nx, ny=ny, Lx=L, f₀=f₀, g=g, H=H, ρ=ρ)
  sol, cl, pr, vs, gr = prob.sol, prob.clock, prob.params, prob.vars, prob.grid

  T = eltype(gr)

  f1 = @. 2cos(k₀*x)*cos(l₀*y)
  f2 = @.  cos(k₀*x+π/10)*cos(2l₀*y)
  f = zeros(dev, T, (gr.nx, gr.ny, nlayers))
  f[:, :, 1] .= f1
  f[:, :, 2] .= f2

  ψtest = zeros(dev, T, size(f))
  MultiLayerQG.set_ψ!(prob, f)
  @. vs.qh = sol
  MultiLayerQG.streamfunctionfrompv!(vs.ψh, vs.qh, pr, gr)
  MultiLayerQG.invtransform!(ψtest, vs.ψh, pr)

  qtest = zeros(dev, T, size(f))
  MultiLayerQG.set_q!(prob, f)
  @. vs.qh = sol
  MultiLayerQG.invtransform!(qtest, vs.qh, pr)

  return isapprox(ψtest, f, rtol=rtol_multilayerqg) && isapprox(qtest, f, rtol=rtol_multilayerqg)
end

"""
    test_paramsconstructor(; kwargs...)

Tests that `Params` constructor works with both mean flow `U` being a floats
(i.e., constant `U` in each layer) or vectors (i.e., `U(y)` in each layer).
"""
function test_mqg_paramsconstructor(dev::Device=CPU(); dt=0.001, stepper="ForwardEuler")
  nx, ny = 32, 34
  L = 2π
  gr = TwoDGrid(dev, nx, L, ny, L)

  nlayers = 2       # these choice of parameters give the
  f₀, g = 1, 1      # desired PV-streamfunction relations
  H = [0.2, 0.8]    # q1 = Δψ1 + 25*(ψ2-ψ1), and
  ρ = [4.0, 5.0]    # q2 = Δψ2 + 25/4*(ψ1-ψ2).

  U1, U2 = 0.1, 0.05

  T = eltype(gr)

  Uvectors = zeros(dev, T, (ny, nlayers))
  Uvectors[:, 1] .= U1
  Uvectors[:, 2] .= U2

  Ufloats = zeros(dev, T, nlayers)
  CUDA.@allowscalar Ufloats[1] = U1
  CUDA.@allowscalar Ufloats[2] = U2

  probUvectors = MultiLayerQG.Problem(nlayers, dev; nx=nx, ny=ny, Lx=L, f₀=f₀, g=g, H=H, ρ=ρ, U=Uvectors)
  probUfloats = MultiLayerQG.Problem(nlayers, dev; nx=nx, ny=ny, Lx=L, f₀=f₀, g=g, H=H, ρ=ρ, U=Ufloats)

  return isapprox(probUfloats.params.U, probUvectors.params.U, rtol=rtol_multilayerqg)
end

function test_mqg_problemtype(dev, T)
  prob1 = MultiLayerQG.Problem(1, dev; T=T)
  prob2 = MultiLayerQG.Problem(2, dev; T=T)

  A = ArrayType(dev)

  return (typeof(prob1.sol)<:A{Complex{T}, 3} && typeof(prob1.grid.Lx)==T && typeof(prob1.vars.u)<:A{T, 3}) && (typeof(prob2.sol)<:A{Complex{T}, 3} && typeof(prob2.grid.Lx)==T && typeof(prob2.vars.u)<:A{T, 3})
end

"""
    test_mqg_rossbywave(; kwargs...)

Evolves a Rossby wave on a beta plane with an imposed zonal flow U and compares
with the analytic solution.
"""
function test_mqg_rossbywave(stepper, dt, nsteps, dev::Device=CPU())
  nlayers = 1
       nx = 64
       Lx = 2π
        β = 2
        U = 0.5

  prob = MultiLayerQG.Problem(nlayers, dev; nx=nx, Lx=Lx, U=U, β=β, stepper=stepper, dt=dt)
  sol, cl, v, p, g = prob.sol, prob.clock, prob.vars, prob.params, prob.grid

  x, y = gridpoints(g)

  # the Rossby wave initial condition
   ampl = 1e-2
  kwave, lwave = 3 * 2π/g.Lx, 2 * 2π/g.Ly
      ω = kwave * ( U - p.β / (kwave^2 + lwave^2) ) # Doppler-shifted Rossby frequency
     q0 = @. ampl * cos(kwave * x) * cos(lwave * y)
     ψ0 = @. - q0 / (kwave^2 + lwave^2)

  MultiLayerQG.set_q!(prob, q0)

  stepforward!(prob, nsteps)
  dealias!(sol, g)
  MultiLayerQG.updatevars!(prob)

  q_theory = @. ampl * cos(kwave * (x - ω / kwave * cl.t)) * cos(lwave * y)

  return isapprox(q_theory, v.q, rtol=g.nx*g.ny*nsteps*1e-12)
end

function test_numberoflayers(dev::Device=CPU())
  prob_nlayers1 = MultiLayerQG.Problem(1, dev)
  prob_nlayers2 = MultiLayerQG.Problem(2, dev)

  return MultiLayerQG.numberoflayers(prob_nlayers1)==1 && MultiLayerQG.numberoflayers(prob_nlayers2)==2
end

function test_mqg_stochasticforcedproblemconstructor(dev::Device=CPU())

  function calcFq!(Fqh, sol, t, clock, vars, params, grid)
    Fqh .= Ffh
    return nothing
  end

  prob = MultiLayerQG.Problem(2, dev; calcFq=calcFq!, stochastic=true)

  return typeof(prob.vars.prevsol)==typeof(prob.sol)
end


using DifferentialEquations
using Sundials
using BenchmarkTools
using Plots

using DifferentialEquations, Plots
function orego(du,u,p,t)
  s,q,w = p
  y1,y2,y3 = u
  du[1] = s*(y2+y1*(1-q*y1-y2))
  du[2] = (y3-(1+y1)*y2)/s
  du[3] = w*(y1-y3)
end
p = [77.27,8.375e-6,0.161]
prob = ODEProblem(orego,[1.0,2.0,3.0],(0.0,360.0),p)
sol = solve(prob)
plot(sol)

plot(sol,vars=(1,2,3))

using BenchmarkTools
prob = ODEProblem(orego,[1.0,2.0,3.0],(0.0,50.0),p)
@btime sol = solve(prob,Tsit5())

@btime sol = solve(prob,Rodas5())

function orego(du,u,p,t)
  s,q,w = p
  y1,y2,y3 = u
  du[1] = s*(y2+y1*(1-q*y1-y2))
  du[2] = (y3-(1+y1)*y2)/s
  du[3] = w*(y1-y3)
end
function g(du,u,p,t)
  du[1] = 0.1u[1]
  du[2] = 0.1u[2]
  du[3] = 0.1u[3]
end
p = [77.27,8.375e-6,0.161]
prob = SDEProblem(orego,g,[1.0,2.0,3.0],(0.0,30.0),p)
sol = solve(prob,SOSRI())
plot(sol)

sol = solve(prob,ImplicitRKMil()); plot(sol)

sol = solve(prob,ImplicitRKMil()); plot(sol)

function orego(du,u,p,t)
  s,q,w = p
  y1,y2,y3 = u
  du[1] = s*(y2+y1*(1-q*y1-y2))
  du[2] = (y3-(1+y1)*y2)/s
  du[3] = w*(y1-y3)
end
p = [60.0,1e-5,0.2]
prob = ODEProblem(orego,[1.0,2.0,3.0],(0.0,30.0),p)
sol = solve(prob,Rodas5(),abstol=1/10^14,reltol=1/10^14)

function onecompartment(du,u,p,t)
  Ka,Ke = p
  du[1] = -Ka*u[1]
  du[2] =  Ka*u[1] - Ke*u[2]
end
p = (Ka=2.268,Ke=0.07398)
prob = ODEProblem(onecompartment,[100.0,0.0],(0.0,90.0),p)

tstops = [24,48,72]
condition(u,t,integrator) = t ∈ tstops
affect!(integrator) = (integrator.u[1] += 100)
cb = DiscreteCallback(condition,affect!)
sol = solve(prob,Tsit5(),callback=cb,tstops=tstops)
plot(sol)

function onecompartment_delay(du,u,h,p,t)
  Ka,Ke,τ = p
  delayed_depot = h(p,t-τ)[1]
  du[1] = -Ka*u[1]
  du[2] =  Ka*delayed_depot - Ke*u[2]
end
p = (Ka=2.268,Ke=0.07398,τ=6.0)
h(p,t) = [0.0,0.0]
prob = DDEProblem(onecompartment_delay,[100.0,0.0],h,(0.0,90.0),p)

tstops = [24,48,72]
condition(u,t,integrator) = t ∈ tstops
affect!(integrator) = (integrator.u[1] += 100)
cb = DiscreteCallback(condition,affect!)
sol = solve(prob,MethodOfSteps(Rosenbrock23()),callback=cb,tstops=tstops)
plot(sol)

p = (Ka = 0.5, Ke = 0.1, τ = 4.0)

function f(du, u, p, t)
    du[1] = -p[1]*u[1] + p[2]*u[2]*u[3]
    du[2] = p[1]*u[1] - p[2]*u[2]*u[3] - p[3]*u[2]*u[2]
    du[3] = u[1] + u[2] + u[3] - 1.
end
M = [1 0 0; 0 1 0; 0 0 0.]
p = [0.04, 10^4, 3e7]
u0 = [1.,0.,0.]
tspan = (0., 1e6)
prob = ODEProblem(ODEFunction(f, mass_matrix = M), u0, tspan, p)
sol = solve(prob, Rodas5())
plot(sol, xscale=:log10, tspan=(1e-6, 1e5), layout=(3,1))

# Robertson Equation DAE Implicit form
function h(out, du, u, p, t)
    out[1] = -p[1]*u[1] + p[2]*u[2]*u[3] - du[1]
    out[2] = p[1]*u[1] - p[2]*u[2]*u[3] - p[3]*u[2]*u[2] - du[2]
    out[3] = u[1] + u[2] + u[3] - 1.
end
p = [0.04, 10^4, 3e7]
du0 = [-0.04, 0.04, 0.0]
u0 = [1.,0.,0.]
tspan = (0., 1e6)
differential_vars = [true, true, false]
prob = DAEProblem(h, du0, u0, tspan, p, differential_vars = differential_vars)
sol = solve(prob, IDA())
plot(sol, xscale=:log10, tspan=(1e-6, 1e5), layout=(3,1))

function f(out, da, a, p, t)
   (L, m, g) = p
   u, v, x, y, T = a
   du, dv, dx, dy, dT = da
   out[1] = x*T/(m*L) - du
   out[2] = y*T/(m*L) - g - dv
   out[3] = u - dx
   out[4] = v - dy
   out[5] = u^2 + v^2 - y*g + T/m
   nothing
end

# Release pendulum from top right
u0 = zeros(5)
u0[3] = 1.0
du0 = zeros(5)
du0[2] = 9.81

p = [1,1,9.8]
tspan = (0.,100.)

differential_vars = [true, true, true, true, false]
prob = DAEProblem(f, du0, u0, tspan, p, differential_vars = differential_vars)
sol = solve(prob, IDA())
plot(sol, vars=(3,4))

function f(out, da, a, p, t)
   L1, m1, L2, m2, g = p

   u1, v1, x1, y1, T1,
   u2, v2, x2, y2, T2 = a

   du1, dv1, dx1, dy1, dT1,
   du2, dv2, dx2, dy2, dT2 = da

   out[1]  = x2*T2/(m2*L2) - du2
   out[2]  = y2*T2/(m2*L2) - g - dv2
   out[3]  = u2 - dx2
   out[4]  = v2 - dy2
   out[5]  = u2^2 + v2^2 -y2*g + T2/m2

   out[6]  = x1*T1/(m1*L1) - x2*T2/(m2*L2) - du1
   out[7]  = y1*T1/(m1*L1) - g - y2*T2/(m2*L2) - dv1
   out[8]  = u1 - dx1
   out[9]  = v1 - dy1
   out[10] = u1^2 + v1^2 + T1/m1 +
                (-x1*x2 - y1*y2)/(m1*L2)*T2 - y1*g
   nothing
end

# Release pendulum from top right
u0 = zeros(10)
u0[3] = 1.0
u0[8] = 1.0
du0 = zeros(10)
du0[2] = 9.8
du0[7] = 9.8

p = [1,1,1,1,9.8]
tspan = (0.,100.)

differential_vars = [true, true, true, true, false,
                     true, true, true, true, false]
prob = DAEProblem(f, du0, u0, tspan, p, differential_vars = differential_vars)
sol = solve(prob, IDA())

plot(sol, vars=(3,4))
plot(sol, vars=(8,9))

using DifferentialEquations, Sundials, Plots

# initial condition
function init_brusselator_2d(xyd)
    N = length(xyd)
    u = zeros(N, N, 2)
    for I in CartesianIndices((N, N))
        x = xyd[I[1]]
        y = xyd[I[2]]
        u[I,1] = 22*(y*(1-y))^(3/2)
        u[I,2] = 27*(x*(1-x))^(3/2)
    end
    u
end

N = 32

xyd_brusselator = range(0,stop=1,length=N)

u0 = vec(init_brusselator_2d(xyd_brusselator))

tspan = (0, 22.)

p = (3.4, 1., 10., xyd_brusselator)

brusselator_f(x, y, t) = ifelse((((x-0.3)^2 + (y-0.6)^2) <= 0.1^2) &&
                                (t >= 1.1), 5., 0.)

using LinearAlgebra, SparseArrays
du = ones(N-1)
D2 = spdiagm(-1 => du, 0=>fill(-2.0, N), 1 => du)
D2[1, N] = D2[N, 1] = 1
D2 = 1/step(xyd_brusselator)^2*D2
tmp = Matrix{Float64}(undef, N, N)
function brusselator_2d_op(du, u, (D2, tmp, p), t)
    A, B, α, xyd = p
    dx = step(xyd)
    N = length(xyd)
    α = α/dx^2
    du = reshape(du, N, N, 2)
    u = reshape(u, N, N, 2)
    @views for i in axes(u, 3)
        ui = u[:, :, i]
        dui = du[:, :, i]
        mul!(tmp, D2, ui)
        mul!(dui, ui, D2')
        dui .+= tmp
    end

    @inbounds begin
        for I in CartesianIndices((N, N))
            x = xyd[I[1]]
            y = xyd[I[2]]
            i = I[1]
            j = I[2]

            du[i,j,1] = α*du[i,j,1] + B + u[i,j,1]^2*u[i,j,2] - (A + 1)*u[i,j,1] + brusselator_f(x, y, t)
            du[i,j,2] = α*du[i,j,2] + A*u[i,j,1] - u[i,j,1]^2*u[i,j,2]
        end
    end
    nothing
end

prob1 = ODEProblem(brusselator_2d_op, u0, tspan, (D2, tmp, p))

sol1 = @time solve(prob1, TRBDF2(autodiff=false));

@gif for t in sol1.t[1]:0.1:sol1.t[end]
    off = N^2
    solt = sol1(t)
    plt1 = surface(reshape(solt[1:off], N, N), zlims=(0, 5), leg=false)
    surface!(plt1, reshape(solt[off+1:end], N, N), zlims=(0, 5), leg=false)
    display(plt1)
end

function brusselator_2d_loop(du, u, p, t)
    A, B, α, xyd = p
    dx = step(xyd)
    N = length(xyd)
    α = α/dx^2
    limit = a -> let N=N
        a == N+1 ? 1 :
        a == 0 ? N :
        a
    end
    II = LinearIndices((N, N, 2))

    @inbounds begin
        for I in CartesianIndices((N, N))
            x = xyd[I[1]]
            y = xyd[I[2]]
            i = I[1]
            j = I[2]
            ip1 = limit(i+1)
            im1 = limit(i-1)
            jp1 = limit(j+1)
            jm1 = limit(j-1)

            ii1 = II[i,j,1]
            ii2 = II[i,j,2]

            du[II[i,j,1]] = α*(u[II[im1,j,1]] + u[II[ip1,j,1]] + u[II[i,jp1,1]] + u[II[i,jm1,1]] - 4u[ii1]) +
            B + u[ii1]^2*u[ii2] - (A + 1)*u[ii1] + brusselator_f(x, y, t)

            du[II[i,j,2]] = α*(u[II[im1,j,2]] + u[II[ip1,j,2]] + u[II[i,jp1,2]] + u[II[i,jm1,2]] - 4u[II[i,j,2]]) +
            A*u[ii1] - u[ii1]^2*u[ii2]
        end
    end
    nothing
end

prob2 = ODEProblem(brusselator_2d_loop, u0, tspan, p)

sol2 = @time solve(prob2, TRBDF2())
sol2_2 = @time solve(prob2, CVODE_BDF())

using SparseDiffTools, SparsityDetection

sparsity_pattern = jacobian_sparsity(brusselator_2d_loop,similar(u0),u0,p,2.0)
jac_sp = sparse(sparsity_pattern)
jac = Float64.(jac_sp)
colors = matrix_colors(jac)
prob3 = ODEProblem(ODEFunction(brusselator_2d_loop, colorvec=colors,jac_prototype=jac_sp), u0, tspan, p)
sol3 = @time solve(prob3, TRBDF2())

using DiffEqOperators
using Sundials
using AlgebraicMultigrid: ruge_stuben, aspreconditioner, smoothed_aggregation
prob6 = ODEProblem(ODEFunction(brusselator_2d_loop, jac_prototype=JacVecOperator{Float64}(brusselator_2d_loop, u0)), u0, tspan, p)
II = Matrix{Float64}(I, N, N)
Op = kron(Matrix{Float64}(I, 2, 2), kron(D2, II) + kron(II, D2))
Wapprox = -I+Op
#ml = ruge_stuben(Wapprox)
ml = smoothed_aggregation(Wapprox)
precond = aspreconditioner(ml)
sol_trbdf2 = @time solve(prob6, TRBDF2(linsolve=LinSolveGMRES())); # no preconditioner
sol_trbdf2 = @time solve(prob6, TRBDF2(linsolve=LinSolveGMRES(Pl=lu(Wapprox)))); # sparse LU
sol_trbdf2 = @time solve(prob6, TRBDF2(linsolve=LinSolveGMRES(Pl=precond))); # AMG
sol_cvodebdf = @time solve(prob2, CVODE_BDF(linear_solver=:GMRES));

function laplacian2d(du, u, p, t)
    A, B, α, xyd = p
    dx = step(xyd)
    N = length(xyd)
    du = reshape(du, N, N, 2)
    u = reshape(u, N, N, 2)
    @inbounds begin
        α = α/dx^2
        limit = a -> let N=N
            a == N+1 ? 1 :
            a == 0 ? N :
            a
        end
        for I in CartesianIndices((N, N))
            x = xyd[I[1]]
            y = xyd[I[2]]
            i = I[1]
            j = I[2]
            ip1 = limit(i+1)
            im1 = limit(i-1)
            jp1 = limit(j+1)
            jm1 = limit(j-1)
            du[i,j,1] = α*(u[im1,j,1] + u[ip1,j,1] + u[i,jp1,1] + u[i,jm1,1] - 4u[i,j,1])
            du[i,j,2] = α*(u[im1,j,2] + u[ip1,j,2] + u[i,jp1,2] + u[i,jm1,2] - 4u[i,j,2])
        end
    end
    nothing
end
function brusselator_reaction(du, u, p, t)
    A, B, α, xyd = p
    dx = step(xyd)
    N = length(xyd)
    du = reshape(du, N, N, 2)
    u = reshape(u, N, N, 2)
    @inbounds begin
        for I in CartesianIndices((N, N))
            x = xyd[I[1]]
            y = xyd[I[2]]
            i = I[1]
            j = I[2]
            du[i,j,1] = B + u[i,j,1]^2*u[i,j,2] - (A + 1)*u[i,j,1] + brusselator_f(x, y, t)
            du[i,j,2] = A*u[i,j,1] - u[i,j,1]^2*u[i,j,2]
        end
    end
    nothing
end
prob7 = SplitODEProblem(laplacian2d, brusselator_reaction, u0, tspan, p)
sol7 = @time solve(prob7, KenCarp4())
M = MatrixFreeOperator((du,u,p)->laplacian2d(du, u, p, 0), (p,), size=(2*N^2, 2*N^2), opnorm=1000)
prob7_2 = SplitODEProblem(M, brusselator_reaction, u0, tspan, p)
sol7_2 = @time solve(prob7_2, ETDRK4(krylov=true), dt=1)
prob7_3 = SplitODEProblem(DiffEqArrayOperator(Op), brusselator_reaction, u0, tspan, p)
sol7_3 = solve(prob7_3, KenCarp4());

using DiffEqDevTools
abstols = 0.1 .^ (5:8)
reltols = 0.1 .^ (1:4)
sol = solve(prob3,CVODE_BDF(linear_solver=:GMRES),abstol=1/10^7,reltol=1/10^10)
test_sol = TestSolution(sol)
probs = [prob2, prob3, prob6]
setups = [Dict(:alg=>CVODE_BDF(),:prob_choice => 1),
          Dict(:alg=>CVODE_BDF(linear_solver=:GMRES), :prob_choice => 1),
          Dict(:alg=>TRBDF2(), :prob_choice => 1),
          Dict(:alg=>TRBDF2(linsolve=LinSolveGMRES(Pl=precond)), :prob_choice => 3),
          Dict(:alg=>TRBDF2(), :prob_choice => 2)
         ]
labels = ["CVODE_BDF (dense)" "CVODE_BDF (GMRES)" "TRBDF2 (dense)" "TRBDF2 (sparse)" "TRBDF2 (GMRES)"]
wp = WorkPrecisionSet(probs,abstols,reltols,setups;appxsol=[test_sol,test_sol,test_sol],save_everystep=false,numruns=3,
  names=labels, print_names=true, seconds=0.5)
plot(wp)

function henon(dz,z,p,t)
  p₁, p₂, q₁, q₂ = z[1], z[2], z[3], z[4]
  dp₁ = -q₁*(1 + 2q₂)
  dp₂ = -q₂-(q₁^2 - q₂^2)
  dq₁ = p₁
  dq₂ = p₂

  dz .= [dp₁, dp₂, dq₁, dq₂]
  return nothing
end

u₀ = [0.1, 0.0, 0.0, 0.5]
prob = ODEProblem(henon, u₀, (0., 1000.))
sol = solve(prob, Vern9(), abstol=1e-14, reltol=1e-14)

plot(sol, vars=[(3,4,1)], tspan=(0,100))

function henon(ddz,dz,z,p,t)
  p₁, p₂ = dz[1], dz[2]
  q₁, q₂ = z[1], z[2]
  ddq₁ = -q₁*(1 + 2q₂)
  ddq₂ = -q₂-(q₁^2 - q₂^2)

  ddz .= [ddq₁, ddq₂]
end

p₀ = u₀[1:2]
q₀ = u₀[3:4]
prob2 = SecondOrderODEProblem(henon, p₀, q₀, (0., 1000.))
sol = solve(prob2, DPRKN6(), abstol=1e-10, reltol=1e-10)

plot(sol, vars=[(3,4)], tspan=(0,100))

H(p, q, params) = 1/2 * (p[1]^2 + p[2]^2) + 1/2 * (q[1]^2 + q[2]^2 + 2q[1]^2 * q[2] - 2/3*q[2]^3)

prob3 = HamiltonianProblem(H, p₀, q₀, (0., 1000.))
sol = solve(prob3, DPRKN6(), abstol=1e-10, reltol=1e-10)

plot(sol, vars=[(3,4)], tspan=(0,100))

function generate_ics(E,n)
  # The hardcoded values bellow can be estimated by looking at the
  # figures in the Henon-Heiles 1964 article
  qrange = range(-0.4, stop = 1.0, length = n)
  prange = range(-0.5, stop = 0.5, length = n)
  z0 = Vector{Vector{typeof(E)}}()
  for q in qrange
    V = H([0,0],[0,q],nothing)
    V ≥ E && continue
    for p in prange
      T = 1/2*p^2
      T + V ≥ E && continue
      z = [√(2(E-V-T)), p, 0, q]
      push!(z0, z)
    end
  end
  return z0
end

z0 = generate_ics(0.125, 10)

function prob_func(prob,i,repeat)
  @. prob.u0 = z0[i]
  prob
end

ensprob = EnsembleProblem(prob, prob_func=prob_func)
sim = solve(ensprob, Vern9(), EnsembleThreads(), trajectories=length(z0))

plot(sim, vars=(3,4), tspan=(0,10))

using DiffEqGPU

function henon_gpu(dz,z,p,t)
  @inbounds begin
    dz[1] = -z[3]*(1 + 2z[4])
    dz[2] = -z[4]-(z[3]^2 - z[4]^2)
    dz[3] = z[1]
    dz[4] = z[2]
  end
  return nothing
end

z0 = generate_ics(0.125f0, 50)
prob_gpu = ODEProblem(henon_gpu, Float32.(u₀), (0.f0, 1000.f0))
ensprob = EnsembleProblem(prob_gpu, prob_func=prob_func)
sim = solve(ensprob, Tsit5(), EnsembleGPUArray(), trajectories=length(z0))

using SciMLTutorials
SciMLTutorials.tutorial_footer(WEAVE_ARGS[:folder],WEAVE_ARGS[:file])


module HHModelCompoenents

function inf_tau_h_rm(V)
    i = (1+exp((V+65)/6))^(-1);
    t = 100*(7*exp((V+60)/11)+10*exp(-(V+60)/25))^(-1)+0.6;
    i, t
end

function inf_tau_m_rm(V)
    i = (1+exp(-(V+38)/7))^(-1);
    t = 10*(5*exp((V+60)/18)+36*exp(-(V+60)/25))^(-1)+0.04;
    i, t
end

function inf_tau_n_rm(V)
    a, b = alpha_beta_n(V);
    i = a / (a+b);
    t = 1 / (a+b);
    i, t
end

function alpha_beta_n(V)
    if V != 60
        a = -0.01*(V+60)/(exp(-(V+60)/10)-1);
    else
        a = 0.1;
    end
    b = 0.125*exp(-(V+70)/80);
    a, b
end

end


"""
Inference module
"""
module Inference

using Distributions

function conv(arr1, arr2)
    arr_conv = zeros(length(arr1)+length(arr2)-1)
    for i in 1:length(arr1)
        for j in 1:length(arr2)
            arr_conv[i+j-1] += arr1[i] * arr2[j]
        end
    end
    arr_conv
end

function fast_conv(arr, n::Int)
    #=  1  2  3
      1 2  3  4
      2 3  4  5
      3 4  5  6
    =#
    if n < 0
        error("Invalid n: $n")
    end
    if n == 0
        return [1.]
    end
    if n == 1
        return arr
    end
    arr2 = conv(arr, arr)
    arr3 = fast_conv(arr2, n ÷ 2)
    if n % 2 == 1
        arr3 = conv(arr3, arr)
    end
    arr3
end

function dice_sum_distribution(base::Int, n::Int)
    # D6 * 2 -> (6, 2) -> 2,3,...,12  1/36,2/36,...
    support = n:(base*n)
    prob = fast_conv(ones(base)/base, n)
    DiscreteNonParametric(support, prob)
end

function prob_fire(firepower::Real, evade::Real, Org_cons::Int)
    f = firepower / 10
    e = evade / 10

    firepower_x = [floor(Int, f), ceil(Int, f)]
    firepower_p = [1+floor(f)-f, f - floor(f)]

    evade_x = [floor(Int, e), ceil(Int, e)]
    evade_p = [1+floor(e)-e, e - floor(e)]

    # (fire_floor, evade_floor), (fire_floor, evade_ceil), (fire_ceil, evade_floor), (fire_ceil, evade_ceil)
    fire_heavy_x = [max(fire_hit - evade_hit, 0) for fire_hit = firepower_x for evade_hit = evade_x]
    fire_light_x = [min(fire_hit, evade_hit) for fire_hit = firepower_x for evade_hit = evade_x]

    fire_hl_p = [fire_p * evade_p for fire_p = firepower_p for evade_p = evade_p]

    max_hit = maximum(fire_heavy_x .+ fire_light_x)
    hit_p = zeros(max_hit+1) # denote prob for 0,1,...,max_hit

    for (hx, lx, p) in zip(fire_heavy_x, fire_light_x, fire_hl_p)
        ha = fast_conv([0.6, 0.4], hx)
        la = fast_conv([0.9, 0.1], lx)
        pp = conv(ha, la)
        hit_p[1:length(pp)] += pp * p
    end

    #Org_cons = 6 # 4

    HP_base_p = ones(2) / 2
    Org_base_p = ones(Org_cons) / Org_cons

    #HP_loss_p = zeros(max_hit*2+1) # denote HP loss for 0, 1, ..., max_hit*2
    #Org_loss_p = zeros(max_hit*Org_cons+1) # denote Org loss for 0, 1, ..., max_hit*Org_cons

    #HP_loss_p[1] = hit_p[1]
    #Org_loss_p[1] = hit_p[1]

    HP_Org_loss_mat = zeros(max_hit*2+1, max_hit*Org_cons+1)
    HP_Org_loss_mat[1, 1] = hit_p[1]

    HP_p = [1.]
    Org_p = [1.]

    for (hit, p) in enumerate(hit_p[2:end])
        HP_p = conv(HP_p, HP_base_p)
        Org_p = conv(Org_p, Org_base_p)

        # marginal
        #HP_loss_p[hit+1:length(HP_p)+hit] += HP_p * p
        #Org_loss_p[hit+1:length(Org_p)+hit] += Org_p * p

        HP_Org_loss_mat[hit+1:length(HP_p)+hit, hit+1:length(Org_p)+hit] += HP_p .* Org_p' * p

    end

    #return HP_loss_p, Org_loss_p
    return HP_Org_loss_mat
end

end


facts("Float split") do
  context("Float32") do
    hi, lo = fsplit(float32(0x8017))
    # assuming non-zero hi, lo
    @fact exponent(hi)-exponent(lo) => greater_than_or_equal(12)
    @fact trailing_zeros(int32(ldexp(significand(hi), 23))) => greater_than_or_equal(12)
  end

  context("Float64") do
    hi, lo = fsplit(float64(0x8000005))
    # assuming non-zero hi, lo
    @fact exponent(hi)-exponent(lo) => greater_than_or_equal(27)
    @fact trailing_zeros(int64(ldexp(significand(hi), 52))) => greater_than_or_equal(27)
  end
end

facts("Error-free multiply") do
  a = float64(0x10000000000001)
  b = 3.0e0
  x, y = err_mul(a, b)
  @fact big(x)+big(y) => big(a)*big(b)
end

facts("Fast error-free multiply") do
  a = float64(0x10000000000001)
  b = 3.0e0
  @fact err_fast_mul(a, b) => err_mul(a, b)
end


function nicoud(L,z;maxiter=10,tol=0.,relax=1.,n_eig_val=3,v0=[],output=true)
    if output
        println("Launching Nicoud...")
        println("Iter    Res:     dz:     z:")
        println("----------------------------------")
    end
    z0=complex(Inf)
    n=0
    flag=itsol_converged

    M=L(1,oplist=["M"],in_or_ex=true) #get mass matrix
    K=L(1,oplist=["K"],in_or_ex=true)
    C=L(1,oplist=["C"],in_or_ex=true)
    d=size(M)[1] #TODO: implement a dimension method for LinearOperatorFamily
    I=SparseArrays.sparse(SparseArrays.I,d,d)
    O=SparseArrays.spzeros(d,d)
    Y=[I O; O M]
    if v0==[]
        v0=ones(ComplexF64,d)
    end
    v0=[v0; z*v0]
    try
        while abs(z-z0)>tol && n< maxiter
            if output; println(n,"\t\t",abs(z-z0),"\t", z );flush(stdout); end
            z0=z
            Q=L(z,oplist=["Q"],in_or_ex=true)
            X=[O -I; K+Q C]
            z,v0=Arpack.eigs(-X,Y,sigma=z0,v0=v0,nev=n_eig_val)
            indexing=sortperm(z.-z0,by=abs)
            z,v0=z[indexing[1]],v0[:,indexing[1]]
            z=z0+relax*(z-z0)
            #TODO consider relaxation in v0
            n+=1
        end
    catch excp
        if output
            println("Error occured:")
            println(excp)
            println(typeof(excp))
            println("...aborted Nicoud!")
        end

        flag=itsol_unknown
        if typeof(excp) <: Arpack.ARPACKException
           flag=itsol_arpack_exception
           if excp==Arpack.ARPACKException(-9999)
               flag=itsol_arpack_9999
           end
        elseif excp== LinearAlgebra.SingularException(0) #This means that the solution is so good that L(z) cannot be LU factorized...TODO: implement other strategy into perturb
          flag=itsol_singular_exception
          L.params[L.eigval]=z
        end
    end
    if flag==itsol_converged
        L.params[L.eigval]=z
        if output;  println(n,"\t\t",abs(z-z0),"\t", z ); end

        if n>=maxiter
            flag=itsol_maxiter
            if output; println("Warning: Maximum number of iterations has been reached!");end

        elseif abs(z-z0)<=tol
            flag=itsol_converged
            if output; println("Solution has converged!"); end
        elseif isnan(z)
            flag=itsol_isnan
            if output; println("Warning: computer arithmetics problem. Eigenvalue is NaN"); end
        else
            if output; println("Warning: This should not be possible....\n If you can read this contact GAM!");end
            flag=itsol_impossible
        end
        if output
            println("...finished Nicoud!")
            println("#####################")
            println(" Nicoud results ")
            println("#####################")
            println("Number of steps: ",n)
            println("Last step parameter variation:",abs(z0-z))
            println("Eigenvalue:",z)
            println("Eigenvalue/(2*pi):",z/2/pi)
        end
    end
    return Solution(L.params,v0[1:d],[],L.eigval), n, flag
end


module Ex7_3c
using FaultDetectionTools
using DescriptorSystems
using LinearAlgebra
using Polynomials
using Test

# Example 7.3c - Nullspace-based synthesis

println("Example 7.3c")

A = diagm([ -2, -3, -1, 2, 3]);
Bu = [1 1 0 0 0]'; Bd = [0 0 2 0 0]';
Bf = [0 0 0 2 2;0 0 0 0 0]';
C = [-1 0 -1 1.5 0; 0 -1 0 0 2.5];
Du = [ 1 1]'; Dd = [1 0]'; Df = [1 0; 1 1];
p = 2; mu = 1; md = 1; mf = 2;       # enter dimensions

# setup the synthesis model
sysf = fdimodset(dss(A,[Bu Bd Bf],C,[Du Dd Df]),c =1:mu, d = mu .+ (1:md), f = (mu+md) .+ (1:mf));

# call of EFDSYN with the options for stability degree -3 and the synthesis
# of a scalar output filter
Qt, Rft, info = efdsyn(sysf, sdeg = -3, rdim = 1);

# compute Q and Rf; scale to meet example
scale = -1
Q = FDFilter(Qt.sys*scale,Qt.outputs,Qt.controls);
Rf = FDFilterIF(Rft.sys*scale,faults = Rft.faults);

# display results
println("Q = $(dss2rm(Q.sys))")
println("Rf = $(dss2rm(Rf.sys))")

end # module


include("iPredict.jl/src/MixtureModelling.jl")
using DelimitedFiles
using Distributions
using StatsPlots
using DataFrames

test_data_1 = rand(LogNormal(1, 0.5), 20)
test_data_2 = rand(LogNormal(3, 0.5), 60)
append!(test_data_1, test_data_2)

relationships = Array{Float64}(undef, length(test_data_1), 2)
mu_arr = [1, 3]
sigma_arr = [0.5, 0.5]
k_arr = [0.5, 0.5]

expectation!(relationships, length(test_data_1), 2, mu_arr, sigma_arr, k_arr, test_data_1)
maximisation_k!(relationships, k_arr)

append!(test_data_1, rand(LogNormal(2, 0.7), 80))

relationships = Array{Float64}(undef, length(test_data_1), 3)
mu_arr = [1, 3, 2]
sigma_arr = [0.5, 0.5, 0.7]
k_arr = [1/3, 1/3, 1/3]

expectation!(relationships, length(test_data_1), 3, mu_arr, sigma_arr, k_arr, test_data_1)
maximisation_k!(relationships, k_arr)


export densitymatrix, ket

"""
    densitymatrix(net, prob, norm=true)

Returns the Density matrix encoded by the neural network `net`, and normalizes
it if `norm==true`.
"""
function densitymatrix(net, prob, norm=true)
    ρ = DenseOperator(basis(prob))
    v = state(prob, net)
    if v isa DiagonalStateWrapper
        v = v.parent
    end
    for i=1:spacedimension(row(v))
        set_index!(row(v), i)
        for j=1:spacedimension(col(v))
            set_index!(col(v),j)
            ρ.data[i,j] = exp(net(v))
        end
    end
    if norm
        ρ.data ./= tr(ρ)
    end
    ρ
end

QuantumOptics.dm(net::NeuralNetwork, prob::AbstractProblem, norm=false) =
    densitymatrix(net, prob, norm)

"""
    ket(net, prob, norm=true)

Returns the state (ket) encoded by the neural network `net`, and normalizes
it if `norm==true`.
"""
function ket(net, prob, norm=true)
    psi = Ket(basis(prob))
    v = state(prob, net)
    for i=1:spacedimension(v)
        set_index!(v, i)
        psi.data[i] = exp(net(v))
    end
    if norm
        normalize!(psi)
    end
    return psi
end


###################################################
# EXEMPLO - resolver PL com a interface CPLEX do Julia
# * Requer CPLEX instalado e CPLEX.jl configurado no Julia
#
# Autor: Leonardo D. Secchin
# Data : 28/06/2021
# leonardosecchin.github.io/topicospo/
###################################################

using CPLEX, JuMP

function cplex_barrier(A, b, c; lb=[], ub=[], eps=1e-8)

    # Lê dimensões
    m,n = size(A)

    # Verifica se as dimensões dos dados são compatíveis
    if (length(b) != m) | (length(c) != n) | (m > n)
        error("Dados com dimensões incompatíveis!")
    end

    # limitantes padrão das variáveis, caso não fornecidos
    if isempty(lb)
        lb = zeros(n,1)
    end
    if isempty(ub)
        ub = 1.0e308 * ones(n,1)
    end

    # estrutura JuMP associada com CPLEX
    P = Model(CPLEX.Optimizer)

    ### configura CPLEX (ver https://www.ibm.com/docs/en/icos/20.1.0?topic=cplex-list-parameters)
    ### os parâmetros são da interface C

    # desliga parte do presolve
    set_optimizer_attribute(P, "CPXPARAM_Preprocessing_Presolve", CPX_OFF)
    set_optimizer_attribute(P, "CPX_PARAM_DEPIND", 0)

    # saídas na tela
    set_optimizer_attribute(P, "CPXPARAM_ScreenOutput", 1)

    # escolhe o algoritmo de barreiras (pontos interiores) do CPLEX
    set_optimizer_attribute(P, "CPXPARAM_LPMethod", CPX_ALG_BARRIER)

    # desliga crossover para Simplex (não purifica a solução para uma solução básica)
    set_optimizer_attribute(P, "CPXPARAM_SolutionType", CPX_NONBASIC_SOLN)

    # eps para convergência (o padrão já é 1e-8)
    set_optimizer_attribute(P, "CPXPARAM_Barrier_ConvergeTol", eps)

    # adiciona variáveis
    @variable(P, x[1:n])

    # define limitantes das variáveis
    for j in 1:n
        set_lower_bound(x[j], lb[j])
        set_upper_bound(x[j], ub[j])
    end

    # adiciona função objetivo
    @objective(P, Min, sum(c[j]*x[j] for j in 1:n))

    # adiciona restrições
    for i in 1:m
        @constraint(P, sum(A[i,j]*x[j] for j in 1:n) == b[i])
    end

    # resolve P
    optimize!(P)

    # retorna solução
    # como capturar número de iterações?
    return getvalue.(x),getobjectivevalue(P),0,((termination_status(P)) == 1 ? 0 : 1)
end


#=
Definitions of the dfe (distribution of fitness effects) functions which define
   selection coefficients in terms of Gamma distributions.
Split off from nn_poplist.jl  on 5/14/18.
=#
using Distributions

function dfe_deleterious( x::Int64; alpha::Float64=0.2, beta::Float64=0.5 )
  global dist_deleterious
  if !isdefined(Main,:dist_deleterious)
    dist_deleterious = Distributions.Gamma(alpha,beta)
  end
  return max(0.1,1.0-rand(dist_deleterious))
end

function dfe_advantageous( x::Int64; alpha::Float64=1.0, beta::Float64=0.5 )
  global dist_advantageous
  if !isdefined(Main,:dist_advantageous)
    dist_advantageous = Distributions.Gamma(alpha,beta)
  end
  return 1.0+rand(dist_advantageous)
end

@doc """ function dfe_mixed( x::Int64; adv_probability::Float64=0.2, alpha_disadv::Float64=0.2, alpha_adv::Float64=1.0, beta_disadv::Float64=1.0, beta_adv::Float64=0.01 )
Mixed advantageous and deleterious.
"""

function dfe_mixed( x::Int64; adv_probability::Float64=0.2, alpha_disadv::Float64=0.2, alpha_adv::Float64=1.0, beta_disadv::Float64=1.0, beta_adv::Float64=0.05 )
  global dist_deleterious
  if !isdefined(Main,:dist_deleterious)
    dist_deleterious = Distributions.Gamma(alpha_disadv,beta_disadv)
  end
  global dist_advantageous
  if !isdefined(Main,:dist_advantageous)
    dist_advantageous = Distributions.Gamma(alpha_adv,beta_adv)
  end
  rand() < adv_probability ? 1.0+rand(dist_advantageous) : max(0.01,1.0-rand(dist_deleterious))
end

@doc """ function dfe_mod( )
Increased fitness for every x which is divisible by modulaus.
The increase in fitness is by a multiple of fit_inc.
"""
function dfe_mod( x::Int64; modulus::Int64=5, fit_inc::Float64=1.2 )
  if x % modulus == 0
	  return fit_inc
  else
    return 1.0  # flat fitness for no selection
  end
end

@doc """function dfe_neutral( x::Int64 )
"""
function dfe_neutral( x::Int64 )
  return 1.0   # flat fitness for no selection
end

@doc """function dfe_fixed( )
"""
function dfe_fixed( x::Int64; s::Float64=0.0 )
  return 1.0 + s   # fitness corresponding to selection coefficient s
end


using Test
using ProjectiveVectors
using LinearAlgebra

function Base.isapprox(x::NTuple{N, <:Number}, y::NTuple{N, <:Number}; kwargs...) where {N}
    isapprox(collect(x), collect(y); kwargs...)
end

@testset "ProjectiveVectors" begin
    @testset "Constructor" begin
        x = PVector([1, 2, 3])
        @test x isa PVector{Int, 1}
        @test data(x) == [1, 2, 3]
        @test dims(x) == (2,)

        x = PVector([1, 2, 3, 4, 5, 6, 7, 8], (2, 2, 1))
        @test x isa PVector{Int, 3}
        @test dims(x) == (2, 2, 1)
        @test PVector([1, 2, 3], [4, 5, 6], [7, 8]) == x

        # broadcasting
        @test x .+ 2 == PVector([1, 2, 3, 4, 5, 6, 7, 8] .+ 2, (2, 2, 1))
        @test complex.(x) isa PVector{Complex{Int64}, 3}
        @test dims(x) == (2,2,1)
    end

    @testset "embed" begin
        z = embed([2, 3])
        @test z isa PVector{Int, 1}
        @test data(z) == [2, 3, 1]

        z = embed([2, 3], [4, 5, 6])
        @test z isa PVector{Int, 2}
        @test data(z) == [2, 3, 1, 4, 5, 6, 1]
        @test dims(z) == (2, 3)

        z = embed([2, 3, 4, 5, 6, 7], (2, 3, 1))
        @test z isa PVector{Int, 3}
        @test dims(z) == (2, 3, 1)
        @test data(z) == [2, 3, 1, 4, 5, 6, 1, 7, 1]
    end

    @testset "Show" begin
        z = PVector([2, 3, 4, 5, 6, 7])
        @test sprint(show, z) == "[2 : 3 : 4 : 5 : 6 : 7]"
        @test sprint(show, z, context=:compact => true) == "[2 : 3 : 4 : 5 : 6 : 7]"

        z = embed([2, 3, 4, 5, 6, 7], (2, 3, 1))
        @test sprint(show, z) == "[2 : 3 : 1] × [4 : 5 : 6 : 1] × [7 : 1]"
        @test sprint(show, z, context=:compact => true) == "[2 : 3 : 1] × [4 : 5 : 6 : 1] × [7 : 1]"
    end

    @testset "Equality" begin
        # all same data layout but should be considered different
        z₁ = PVector([2, 3, 4, 5, 6, 7])
        z₂ = PVector([2, 3, 4, 5, 6, 7], (2, 2))
        z₃ = PVector([2, 3, 4, 5, 6, 7], (3, 1))

        @test z₁ ≠ z₂
        @test z₁ ≠ z₃
        @test z₂ ≠ z₃
        @test z₂ == z₂

        z₂_2 = PVector([2, 3, 4, 5.0, 6.0, 7.0], (2, 2))
        @test z₂ == z₂_2
    end

    @testset "Linear Algebra" begin
        z = embed([2.0, 3, 4, 5, 6, 7], (2, 3, 1))
        @test norm(z) ≈ (3.7416573867739413, 8.831760866327848, 7.0710678118654755)
        @test norm(normalize(z)) ≈ (1.0, 1.0, 1.0)
        @test norm(normalize(z, 2), 2) ≈ (1.0, 1.0, 1.0)
        @test norm(normalize(z, Inf), Inf) ≈ (1.0, 1.0, 1.0)

        z = embed([2.0, 3, 4, 5, 6, 7])
        @test norm(z) ≈ (11.832159566199232,)
        @test norm(normalize(z)) ≈ (1.0,)

        z = embed(rand(ComplexF64, 6), (2, 3, 1))
        @test norm(normalize(z)) ≈ (1.0, 1.0, 1.0)
        @test norm(normalize(z, 2), 2) ≈ (1.0, 1.0, 1.0)
        @test norm(normalize(z, Inf), Inf) ≈ (1.0, 1.0, 1.0)
    end

    @testset "affine_chart" begin
        w = [2.0, 3, 4, 5, 6, 7]
        z = embed(w, (2, 3, 1), normalize=true)
        @test affine_chart(z) ≈ w
        @test affine_chart!(zeros(6), z) ≈ w

        z = normalize(embed(w))
        @test affine_chart(z) ≈ w
        @test affine_chart!(zeros(6), z) ≈ w

        z = embed(w, normalize=true)
        @test affine_chart(z) ≈ w
        @test affine_chart!(zeros(6), z) ≈ w

        for T = [Float64, ComplexF64]
            w = rand(ComplexF64, 6)
            z = embed(w, (2, 3, 1), normalize=true)
            @test norm_affine_chart(z, 2) ≈ norm(w, 2)
            @test norm_affine_chart(z, Inf) ≈ norm(w, Inf)
        end
    end

    @testset "Dot, Fubini-Study" begin
        z = embed(rand(ComplexF64, 6), normalize=true)
        @test abs.(dot(z, z)) ≈ (1.0,)
        z = embed(rand(ComplexF64, 6), (2, 4), normalize=true)
        @test abs.(dot(z, z)) ≈ (1.0, 1.0)
        @test fubini_study(z, z) ≈ 0.0
    end

    @testset "Componentes" begin
        v = PVector([1,2,3])
        w = PVector([4, 5])

        @test combine(v, w) == PVector([1,2,3], [4, 5])
        @test v × w == PVector([1,2,3], [4, 5])
        @test v × w × v == combine(v, w, v) == PVector([1,2,3], [4, 5], [1, 2, 3])

        a, b, c = components(v × w × v)
        @test a == c == v
        @test b == w

        @test component(v × w × v, 2) == w
        @test component(v × w × v, 3) == v
        @test (v × w × v)[2,:] == w
        @test (v × w × v)[3,:] == v
    end
end


export moment, moment_value, monomial

struct Moment{T, MT <: AbstractMonomial} <: AbstractMoment{T}
    α::T
    x::MT
end

"""
    moment(α, m::AbstractMonomial)

Creates the moment of the monomial `m` of value `α`.
"""
moment(α, m::AbstractMonomial) = Moment(α, m)

"""
    moment_value(m::AbstractMomentLike)

Returns the value of the moment `m`.

## Examples

Calling `moment_value(moment(3.1, x*y^2))` should return `3.1`.
"""
moment_value(m::Moment) = m.α

"""
    monomial(m::AbstractMomentLike)

Returns the monomial of the moment `m`.

## Examples

Calling `monomial(moment(3.1, x*y^2))` should return `x*y^2`.
"""
MP.monomial(m::Moment) = m.x

for f in [:variables, :nvariables, :exponents, :degree, :powers]
    @eval begin
        MP.$f(m::AbstractMomentLike) = $f(monomial(m))
    end
end