using ApproxFun, SingularIntegralEquations, Base.Test

println("Hilbert test")

x=Fun(identity)
f=(exp(x)/(sqrt(1-x)*sqrt(x+1)))
@test_approx_eq (Hilbert(f|>space,0)*f)(.1) (-0.8545003781055088)
@test_approx_eq (Hilbert(0)*f)(.1) (-0.8545003781055088)
@test_approx_eq (Hilbert()*f)(.1) 1.1404096104609646386

x=Fun(identity,[-1,2])
f=(exp(x)/(sqrt(2-x)*sqrt(x+1)))
@test_approx_eq (Hilbert(f|>space,0)*f)(.1) 0.49127801561694168644
@test_approx_eq (Hilbert(0)*f)(.1) 0.49127801561694168644
@test_approx_eq (Hilbert()*f)(.1) 1.6649936695644078289

x=Fun(identity)
f=(exp(x)*(sqrt(1-x)*sqrt(x+1)))
@test_approx_eq (Hilbert()*f)(.1) 0.43723982258866913063

x=Fun(identity,[-1,2])
f=(exp(x)*(sqrt(2-x)*sqrt(x+1)))
@test_approx_eq (Hilbert()*f)(.1) 2.1380903070701673244

x=Fun(identity)
d=domain(x)
w=1/sqrt(1-x^2)
H=Hilbert(d)
B=ldirichlet(d)

for a in [sqrt(sqrt(5)-2)/2,1.,10.]
    L=H[w]+1/a/sqrt(1+a^2)*x
    u=[B,L]\[1.]
    usol = (1+a^2)/(x^2+a^2)
    @test norm(u-usol) <= eps(1000/a)
end

x = Fun(identity)
w = 1/sqrt(1-x^2)
H = Hilbert(space(w))
@test_approx_eq  (H[w]*exp(x))(.1) hilbert(w*exp(x))(.1)

x = Fun(identity)
w = sqrt(1-x^2)
H = Hilbert(space(w))
@test_approx_eq (H[w]*exp(x))(.1) hilbert(w*exp(x))(.1)

println("Stieltjes test")

ds1 = JacobiWeight(-.5,-.5,ApproxFun.ChebyshevDirichlet{1,1}())
ds2 = JacobiWeight(-.5,-.5,Chebyshev())
rs = Chebyshev([2.,4.+3im])
f1 = Fun(x->exp(x)/sqrt(1-x^2),ds1)
f2 = Fun(x->exp(x)/sqrt(1-x^2),ds2)
S = Stieltjes(ds1,rs)

z = 3.+1.5im
@test_approx_eq (S*f1)(z) stieltjes(f2,z) #val,err = quadgk(x->f1(x)./(z-x),-1.,1.)
# Operator 1.1589646343327578 - 0.7273679005911196im
# Function 1.1589646343327455 - 0.7273679005911283im

ds1 = JacobiWeight(.5,.5,Ultraspherical(1))
ds2 = JacobiWeight(.5,.5,Chebyshev())
rs = Chebyshev([2.,4.])
f1 = Fun(x->exp(x)*sqrt(1-x^2),ds1)
f2 = Fun(x->exp(x)*sqrt(1-x^2),ds2)
S = Stieltjes(ds1,rs)

z = 3.
@test_approx_eq (S*f1)(z) stieltjes(f2,z) #val,err = quadgk(x->f1(x)./(z-x),-1.,1.;reltol=eps())
# Operator 0.6616422557285478 + 0.0im
# Function 0.661642255728541 - 0.0im

println("Stieltjes integral test")

ds1 = JacobiWeight(-.5,-.5,ApproxFun.ChebyshevDirichlet{1,1}())
ds2 = JacobiWeight(-.5,-.5,Chebyshev())
rs = Chebyshev([2.,4.])
f1 = Fun(x->exp(x)/sqrt(1-x^2),ds1)
f2 = Fun(x->exp(x)/sqrt(1-x^2),ds2)
S = Stieltjes(ds1,rs,0)

z = 3.
@test_approx_eq (S*f1)(z) SingularIntegralEquations.stieltjesintegral(f2,z)
# Operator 3.6322473044237698 + 0.0im
# Function 3.6322473044237515

ds1 = JacobiWeight(.5,.5,Ultraspherical(1))
ds2 = JacobiWeight(.5,.5,Chebyshev())
rs = Chebyshev([2.0,4.0])
f1 = Fun(x->exp(x)*sqrt(1-x^2),ds1)
f2 = Fun(x->exp(x)*sqrt(1-x^2),ds2)
S = Stieltjes(ds1,rs,0)

z = 3.0
@test_approx_eq (S*f1)(z) SingularIntegralEquations.stieltjesintegral(f2,z)
# Operator 1.7772163062194861 + 0.0im
# Function 1.7772163062194637

println("Cauchy test")

f2=Fun(sech,PeriodicLine())
@test_approx_eq cauchy(f2,1.+im) (0.23294739894134472 + 0.10998776661109881im )
@test_approx_eq cauchy(f2,1.-im) (-0.23294739894134472 + 0.10998776661109881im )

f=Fun(sech,Line())
@test_approx_eq cauchy(f,1.+im) cauchy(f2,1.+im)

f=Fun(z->exp(exp(0.1im)*z+1/z),Laurent(Circle()))

@test_approx_eq hilbert(f,exp(0.2im)) hilbert(Fun(f,Fourier),exp(0.2im))
@test_approx_eq hilbert(f,exp(0.2im)) -hilbert(reverseorientation(f),exp(0.2im))
@test_approx_eq hilbert(f,exp(0.2im)) -hilbert(reverseorientation(Fun(f,Fourier)),exp(0.2im))

@test_approx_eq cauchy(f,0.5exp(0.2im)) cauchy(Fun(f,Fourier),0.5exp(0.2im))
@test_approx_eq cauchy(f,0.5exp(0.2im)) -cauchy(reverseorientation(f),0.5exp(0.2im))
@test_approx_eq cauchy(f,0.5exp(0.2im)) -cauchy(reverseorientation(Fun(f,Fourier)),0.5exp(0.2im))
@test_approx_eq cauchy(f,0.5exp(0.2im)) (OffHilbert(space(f),Laurent(Circle(0.5)))*f)(0.5exp(0.2im))/(2im)

f=Fun(z->exp(exp(0.1im)*z+1/(z-1.)),Laurent(Circle(1.,0.5)))
@test_approx_eq cauchy(f,0.5exp(0.2im)) cauchy(Fun(f,Fourier),0.5exp(0.2im))
@test_approx_eq cauchy(f,0.5exp(0.2im)) -cauchy(reverseorientation(f),0.5exp(0.2im))
@test_approx_eq cauchy(f,0.5exp(0.2im)) -cauchy(reverseorientation(Fun(f,Fourier)),0.5exp(0.2im))

Γ=Circle()∪Circle(0.5)
f=depiece([Fun(z->z^(-1),Γ[1]),Fun(z->z,Γ[2])])
A=I-(f-Fun(one,space(f)))*Cauchy(-1)
u=A\(f-Fun(one,space(f)))

@test_approx_eq 1+cauchy(u,.1) 1
@test_approx_eq 1+cauchy(u,.8) 1/0.8
@test_approx_eq 1+cauchy(u,2.) 1

c1=0.5+0.1;r1=3.;
c2=-0.1+.2im;r2=0.3;
d1=Circle(c1,r1)
d2=Circle(c2,r2)
z=Fun(identity,d2);
C=Cauchy(Space(d1),Space(d2))
@test norm((C*Fun(exp,d1)-Fun(exp,d2)).coefficients)<100eps()

C2=Cauchy(Space(d2),Space(d1))
@test norm((C2*Fun(z->exp(1/z)-1,d2)+Fun(z->exp(1/z)-1,d1)).coefficients)<100000eps()

c1=0.1+.1im;r1=.4;
c2=-2.+.2im;r2=0.3;
d1=Circle(c1,r1)
d2=Circle(c2,r2)
@test norm((Cauchy(d1,d2)*Fun(z->exp(1/z)-1,d1)+Fun(z->exp(1/z)-1,d2)).coefficients)<2000eps()

# complex contour

#Legendre uses FastGuassQuadrature
f=Fun(exp,Legendre())
#@test_approx_eq cauchy(f,.1+0.000000000001im) cauchy(f,.1,+)
#@test_approx_eq cauchy(f,.1-0.000000000001im) cauchy(f,.1,-)
#@test_approx_eq (cauchy(f,.1,+)-cauchy(f,.1,-)) exp(.1)

ω=2.
d=Interval(0.5im,30.0im/ω)
x=Fun(identity,Legendre(d))
@test_approx_eq cauchy(exp(im*ω*x),1.+im) (-0.025430235512791915911 + 0.0016246822285867573678im)

println("Arc test")

a=Arc(0.,1.,0.,π/2)
ζ=Fun(identity,a)
f=Fun(exp,a)*sqrt(abs((ζ-1)*(ζ-im)))
z=.1+.2im
#@test_approx_eq cauchy(f,z) sum(f/(ζ-z))/(2π*im)
z=exp(.1im)
@test_approx_eq hilbert(f,z) im*(cauchy(f,z,+)+cauchy(f,z,-))

println("Functional test")
z=.1+.2im
x=Fun(identity)
f=exp(x)*sqrt(1-x^2)
@test_approx_eq Stieltjes(space(f),z)*f stieltjes(f,z)

a=Arc(0.,1.,0.,π/2)
ζ=Fun(identity,a)
f=Fun(exp,a)*sqrt(abs((ζ-1)*(ζ-im)))
H=Hilbert()
z=exp(.1im)
@test_approx_eq (H*f)(z) hilbert(f,z)

println("Logkernel test")

a=1.0;b=2.0
d=Interval(a,b)
z=Fun(d)
f=real(exp(z)/(sqrt(z-a)*sqrt(b-z)))
S=space(f)
x=4.0+2im;
@test_approx_eq linesum(f*log(abs(x-z))) logkernel(f,x)*π

a=1.0;b=2.0+im
d=Interval(a,b)
z=Fun(d)
f=real(exp(z)/(sqrt(z-a)*sqrt(b-z)))
S=space(f)
x=4.0+2im;
@test_approx_eq linesum(f*log(abs(x-z))) 13.740676344264614
@test_approx_eq linesum(f*log(abs(x-z))) logkernel(f,x)*π

linesum(f)
sum(f)
a=1.0;b=2.0+im
d=Interval(a,b)
z=Fun(d)
f=real(exp(z)*(sqrt(z-a)*sqrt(b-z)))
S=space(f)
x=4.0+2im;
@test_approx_eq linesum(f*log(abs(x-z))) logkernel(f,x)*π

a=1.0;b=2.0
d=Interval(a,b)
z=Fun(d)
f=real(exp(z)/(sqrt(z-a)*sqrt(b-z)))
x=1.5
@test_approx_eq (SingularIntegral(space(f),0)*f)(x) logkernel(f,x)

f=real(exp(z)*(sqrt(z-a)*sqrt(b-z)))
x=1.5
@test_approx_eq (SingularIntegral(space(f),0)*f)(x) logkernel(f,x)

a=1.0;b=2.0+im
d=Interval(a,b)
z=Fun(d)
f=real(exp(z)/(sqrt(z-a)*sqrt(b-z)))
x=1.5+0.5im
H=SingularIntegral(space(f),0)
@test_approx_eq (H*f)(x) logkernel(f,x)

f=real(exp(z)*(sqrt(z-a)*sqrt(b-z)))
@test_approx_eq (SingularIntegral(space(f),0)*f)(x) logkernel(f,x)

a=1.0;b=2.0+im
d=Interval(a,b)
z=Fun(d)
f=real(exp(z)/(sqrt(z-a)*sqrt(b-z)))
S=JacobiWeight(-0.5,-0.5,ChebyshevDirichlet{1,1}(d))
H=OffSingularIntegral(S,Chebyshev([3,4]),0)
@test_approx_eq (H*f)(3.5) logkernel(f,3.5)

H=OffSingularIntegral(S,Chebyshev([3,4.0+im]),0)
@test_approx_eq (H*f)(3.5+0.5im) logkernel(f,3.5+0.5im)

## Circle

d=Circle(0.2,3.0)
S=Fourier(d)
ζ=Fun(d)
f=real(ζ+1/(ζ-0.1))
#z=0.1+0.1im;@test_approx_eq linesum(log(abs(ζ-z))*f) logkernel(f,z)*π
#z=5.0+0.1im;@test_approx_eq linesum(log(abs(ζ-z))*f) logkernel(f,z)*π

d=Circle(0.2,3.0)
S=Fourier(d)
H=Hilbert(S,0)
ζ=Fun(d)
f=real(ζ+1/(ζ-0.1))
z=0.2+3im;@test_approx_eq (H*f)(z) logkernel(f,z)

if isdir(Pkg.dir("FastGaussQuadrature"))
    println("Stieltjes moment test")
    include("stieltjesmomentTest.jl")
end

println("Curve Test")

include("CurveTest.jl")

println("FundamentalSolutions test")

include("FundamentalSolutionsTest.jl")

println("Convolution ProductFun test")

include("convolutionProductFunTest.jl")

println("LowRankMatrix test")

include("LowRankMatrixTest.jl")

println("HierarchicalVector test")

include("HierarchicalVectorTest.jl")

println("Hierarchical solve test")

include("hierarchicalsolveTest.jl")

# Testing Travis CI abilities. Currently, Examples tests are slow.
#println("\nExamples test\n")
#
#include("ExamplesTest.jl")

println("WienerHopfTest")
include("WienerHopfTest.jl")


function solve(
    i::MinimumBudget{UniformMatroidInstance{Float64, Maximise}, Int},
    ::DynamicProgramming,
)
    # Consider items above δ.
    # Recurrence relation:
    #  V[m, t, δ] = max { V[m, t, δ + 1],  v_δ + V[m - 1, t - w_δ, δ + 1] }
    #                     \_ no item δ _/ \________ take item δ ________/
    V = Array{Float64, 3}(
        undef,
        i.instance.m,
        dimension(i) + 1,
        i.min_budget + 1,
    )
    S = Dict{Tuple{Int, Int, Int}, Vector{Int}}()

    # Initialise: µ == 1, just take the best element among [δ+1, d], as long as
    # the budget constraint is satisfied; δ == d, take nothing, whatever the
    # budget; β == 0, a simple m-set problem.
    for µ in 1:(i.instance.m) # β == 0
        for δ in (dimension(i) - 1):-1:0 # δ < d
            # Due to δ, there may be too few items (with respect to µ).
            # i.e., there are d - δ items to consider.
            m_ = min(dimension(i) - δ, µ)

            items =
                collect(
                    partialsortperm(
                        i.instance.rewards[(δ + 1):dimension(i)],
                        1:m_,
                        rev=true,
                    ),
                ) .+ δ
            V[µ, δ + 1, 0 + 1] = sum(i.instance.rewards[o] for o in items)
            S[µ, δ, 0] = items
        end
    end

    for β in 0:(i.min_budget)
        # δ == d
        for µ in 1:(i.instance.m)
            if β == 0
                # Optimum solution: take no object.
                V[µ, dimension(i) + 1, β + 1] = 0.0
                S[µ, dimension(i), β] = Int[]
            else
                # No solution.
                V[µ, dimension(i) + 1, β + 1] = -Inf
                S[µ, dimension(i), β] = [-1] # Nonsensical solution.
            end
        end

        # δ < d
        for δ in (dimension(i) - 1):-1:0
            # Filter the available objects (i.e. above δ+1 and satisfying the
            # budget constraint) and retrieve just their indices.
            all_objects = collect(enumerate(i.weights))
            all_objects = collect(filter(o -> o[1] >= δ + 1, all_objects))
            all_objects = collect(filter(o -> o[2] >= β, all_objects))
            all_objects = collect(t[1] for t in all_objects)

            # No feasible solution?
            if length(all_objects) == 0
                V[1, δ + 1, β + 1] = -Inf
                S[1, δ, β] = [-1] # Nonsensical solution.
                continue
            end

            # Take the best one available object.
            v, x = findmax(collect(i.instance.rewards[t] for t in all_objects))
            x = all_objects[x]

            V[1, δ + 1, β + 1] = v
            S[1, δ, β] = [x]
        end
    end

    # Dynamic part.
    for β in 1:(i.min_budget)
        for µ in 2:(i.instance.m)
            for δ in (dimension(i) - 1):-1:0
                remaining_budget_with_δ = max(0, β - i.weights[δ + 1])
                take_δ =
                    i.instance.rewards[δ + 1] +
                    V[µ - 1, δ + 1 + 1, remaining_budget_with_δ + 1]
                dont_take_δ = V[µ, δ + 1 + 1, β + 1]

                both_subproblems_infeasible =
                    -Inf == take_δ && -Inf == dont_take_δ

                if both_subproblems_infeasible
                    V[µ, δ + 1, β + 1] = -Inf
                    S[µ, δ, β] = [-1]
                elseif take_δ >= dont_take_δ
                    V[µ, δ + 1, β + 1] = take_δ
                    S[µ, δ, β] =
                        vcat(δ + 1, S[µ - 1, δ + 1, remaining_budget_with_δ])
                else
                    V[µ, δ + 1, β + 1] = dont_take_δ
                    S[µ, δ, β] = S[µ, δ + 1, β]
                end
            end
        end
    end

    S_refined = Dict(β => S[i.instance.m, 0, β] for β in 0:(i.min_budget))
    return MinBudgetedUniformMatroidSolution(
        i,
        S[i.instance.m, 0, i.min_budget],
        S_refined,
    )
end


using Statistics
using Flux: onehotbatch, onecold, crossentropy, throttle
using Base.Iterators: repeated
#using CuArrays if you want to use GPU

imgs = Flux.Data.MNIST.images()
labels = Flux.Data.MNIST.labels();

X = hcat(float.(reshape.(imgs, :))...) #stack all the images
Y = onehotbatch(labels, 0:9) # just a common way to encode categorical variables

m = Chain(
  Dense(28^2, 32, relu),
  Dense(32, 10),
  softmax)

loss(x, y) = crossentropy(m(x), y)
opt = ADAM(); # popular stochastic gradient descent variant

accuracy(x, y) = mean(onecold(m(x)) .== onecold(y)) # cute way to find average of correct guesses

dataset = repeated((X,Y),200) # repeat the data set 200 times, as opposed to @epochs 200 ...
evalcb = () -> @show(loss(X, Y)) # callback to show loss

Flux.train!(loss, params(m), dataset, opt, cb = throttle(evalcb, 10)); #took me ~5 minutes to train on CPU

# Same preprocessing
test_X = hcat(float.(reshape.(Flux.Data.MNIST.images(:test), :))...)
test_Y = onehotbatch(Flux.Data.MNIST.labels(:test), 0:9);

m(test_X[:,5287]) # Note the 7th index ( corresponding to the digit 6 ) is nearly 1

#decode
onecold(m(test_X[:,5287])) - 1 #minus 1 since we start from 0, but indexing in Julia starts at 1

# Training set accuracy
accuracy(X, Y)

#######
using Flux, Metalhead
using Flux: @epochs
resnet = ResNet().layers

model = Chain(
  resnet[1:end-2],
  x -> reshape(x, 512),
  Dense(512, 13)
)

include("dataloader.jl")

dataset = [gpu.(load_batch(13*5)) for i in 1:10]

model = model |> gpu
opt = ADAM()

loss(x,y) = Flux.Losses.logitcrossentropy(model(x), y)

ps = Flux.params(model[2:end])

@epochs 2 Flux.train!(loss, ps, dataset, opt)

imgs, labels = gpu.(load_batch(13))
display(model(imgs))

labels

using Statistics
using Flux: onehotbatch, onecold, crossentropy, throttle
using Base.Iterators: repeated

accuracy(x, y) = mean(onecold(model(x)) .== onecold(y)) # cute way to find average of correct guesses

X = imgs
Y = labels

X1, Y1 = dataset[1][1], dataset[1][2]
X2, Y2 = dataset[2][1], dataset[2][2]

accuracy(X, Y)
accuracy(X1, Y1)
accuracy(X2, Y2)


# wls object is not actively maintained

function calcWlsCoefficients(obj::WlsObject, inPt::Point)
  m = length(obj.b)
  datalen = size(obj.inputs, 1)
  firstTerm = zeros(m, m)
  secondTerm = zeros(m)

  for p in 1:datalen
    curPt = obj.inputs[p, :]
    w = obj.weightFunc(norm(inPt - curPt), obj.EPS)
    for i in 1:m
      for j in 1:m
        firstTerm[i, j] += w * obj.matrix[i, j](obj.vars => curPt)
      end
      secondTerm[i] += w * obj.b[i](obj.vars => curPt) * obj.outputs[p]
    end
  end

  coefficients = firstTerm \ secondTerm
  return coefficients
end

"""
    `(wo::WlsObject)(inputPoint::Point, value::Bool)`

Approximates the `WlsObject` at `inputPoint`.
"""
function (obj::WlsObject)(inPt::Point)
  c = calcWlsCoefficients(obj, inPt)
  poly = polynomial(c, obj.b)
  return poly(obj.vars => inPt)
end

function plotWls(obj::WlsObject)
  inputDim = size(obj.inputs, 2)
  if inputDim < 1 || inputDim > 2
    error("Cannot plot $inputDim dimensional graph")
  end
  x = [linspace(-2, 2, 21) for i in 1:inputDim]
  println(x)
  println(length(x))
  if inputDim == 2
    return plot(x[1], x[2], [obj([a, b], true) for b in x[1], a in x[2]], linetype=:surface)
  else
    return plot(x[1], [obj([a], true) for a in x[1]])
  end
end


include("../src/rv_model_keplerian_simple.jl")

if !isdefined(:PDMats) using PDMats  end # Positive Definite Matrices
if !isdefined(:Optim)  using Optim   end # Optimizing over non-linear parameters
if !isdefined(:QuadGK) using QuadGK  end # 1-d integration
if !isdefined(:Cuba)   using Cuba    end # multi-d integration

# Read dataset
data = readdlm("../data/rvs_0001.txt");
times = data[:,1]
vels = data[:,2]
sigma_obs = data[:,3]

Sigma0 = PDMat(make_Sigma(times,sigma_obs,0.0))
SigmaBig = PDMat(make_Sigma(times,sigma_obs,std(vels)))
param0 = [0.]
res = laplace_approximation_kepler_model(param0,times,vels,SigmaBig,loglikelihood_fixed_jitter)
param0 = res[2].minimizer;
laplace_approximation_kepler_model_jitter_separate(param0,times,vels,Sigma0)[1]/log(10)

param1a = [42.0,5.0,0.01,0.01,1.,0.]
res = laplace_approximation_kepler_model(param1a,times,vels,SigmaBig,loglikelihood_fixed_jitter)
param1a = res[2].minimizer;
laplace_approximation_kepler_model_jitter_separate(param1a,times,vels,Sigma0)[1]/log(10)

param1b = [12.1,3.0,0.01,0.01,1.,0.]
res = laplace_approximation_kepler_model(param1b,times,vels,SigmaBig,loglikelihood_fixed_jitter)
param1b = res[2].minimizer;
laplace_approximation_kepler_model_jitter_separate(param1b,times,vels,Sigma0)[1]/log(10)

#param2 = [42.0,5.0,0.01,0.01,1.,12.1,3.0,0.01,0.01,1.,0.]
param2 = vcat(param1a[1:5], param1b[1:5], 0.0)
res = laplace_approximation_kepler_model(param2,times,vels,SigmaBig,loglikelihood_fixed_jitter)
param2 = res[2].minimizer;
laplace_approximation_kepler_model_jitter_separate(param2,times,vels,Sigma0)[1]/log(10)


using Optim
using Cubature
using Roots
using StatsFuns

# import Statistics: median, pdf, cdf
"pre-process pdf and cdf, given fixed pdf and cdf at x0, compute estimated support and check if pdf is proper"
function pre_process(x0::Array{T,2}, Fx0::Array{T,2}, pdf::Function, cdf::Function, dpdf::Function, quantbound::Function; yshift = 0.) where T<:Float64
    quantbound_fixed(p) = quantbound(x0, Fx0, p)
    pdf_fixed(y) = pdf(x0, Fx0, y)
    cdf_fixed(y) = cdf(x0, Fx0, y)
    dpdf_fixed(y) = dpdf(x0, Fx0, y)
    support = [1e-2, 5.]
    function support_comp!(pdf, support)
      current = pdf(support[1])
      for i in 1:5
        next = pdf(support[1]/10)
        if next < current
          support[1] /= 10
        else
          break
        end
        current = next
      end
      while pdf(support[2]) > 1e-6
        support[2] *= 1.2
      end
      support[2] += yshift
      # should make sure CDF(support[2]) - CDF(support[1]) > .96 to make 95% CI possible
      INT = cdf_fixed(support[2]) - cdf_fixed(support[1])
      if INT < .96
        @warn "pdf integral $INT"
      end
      return INT
    end
    INT = support_comp!(pdf_fixed, support)
    return (pdf_fixed, cdf_fixed, dpdf_fixed, quantbound_fixed, support, INT)
end

"wrap up all statistics computation"
function summary_comp(pdf_fixed::Function, cdf_fixed::Function, dpdf_fixed::Function, quantbound_fixed::Function, support::Array{T,1};
                       px = .5, confidence_level = .95) where T<:Float64
    quant_p, error_quant = quantile(cdf_fixed, quantbound, support; p=px)
    med, error_med = median(cdf_fixed, quantbound, support)
    mod = mode(pdf_fixed, support)
    CI_equal, error_CI_eq = credible_interval(cdf_fixed, quantbound, support;
                                                mode=:equal, wp = confidence_level)
    # CI_narrow, error_CI_nr = credible_interval(cdf_fixed, quantbound, support;
    #                                             mode=:narrow, wp = confidence_level)
    quantileInfo = (level = px, value = quant_p, error = error_quant)
    medianInfo = (value = med, error = error_med)
    CIequalInfo = (equal = CI_equal, error = error_CI_eq)
    # CInarrowInfo = (equal = CI_narrow, error = error_CI_nr)
    CInarrowInfo = nothing
    DistributionInfo = (quantile = quantileInfo, median = medianInfo, mode = mod, CIequal = CIequalInfo, CInarrow = CInarrowInfo)
    return DistributionInfo
end

"""
Given pdf, cdf and maybe pdf_deriv,
compute median, quantile, mode, symmetric/narrowest credible interval.
Warning: only for normalized values
"""
function median(cdf::Function, quantbound::Function, support::Array{T,1}; pdf = nothing, pdf_deriv=nothing) where T<:Float64
    med, err = quantile(cdf, quantbound, support)
    return med, err, bound
end

function quantile(cdf::Function, quantbound::Function, support::Array{T,1}; pdf = nothing, pdf_deriv=nothing, p::T=.5) where T<:Float64
    bound = support
    # try
    #   bound = quantbound(p)
    # catch err
    # end
    quant = fzero(y0 -> cdf(y0) - p, bound[1], bound[2])
    err = abs(p-cdf(quant))/p
    # status = err < 1e-5 ? 1 : 0
    return quant, err, bound
end

function mode(pdf::Function, support::Array{T,1}; cdf = nothing, pdf_deriv=nothing) where T<:Float64
    # maximize the pdf
    routine = optimize(x -> -pdf(x), support[1], support[2])
    mod = Optim.minimizer(routine)
    return mod
end

function credible_interval(cdf::Function, quantbound::Function, support::Array{T,1};
                            pdf=nothing, pdf_deriv=nothing, wp::T=.95, mode=:equal) where T<:Float64
    return credible_interval(cdf, quantbound, support, Val(mode); pdf_deriv=pdf_deriv, wp=wp)
end

function credible_interval(cdf::Function, quantbound::Function, support::Array{T,1}, ::Val{:equal};
                            pdf=nothing, pdf_deriv=nothing, wp::T=.95) where T<:Float64
    lower_qp = (1 - wp) / 2
    upper_qp = 1 - lower_qp
    lower_quant = lower_qp # random initialization
    try
        lower_quant = quantile(cdf, quantbound, support; p=lower_qp)[1]
    catch err
        wp = .9
        lower_qp = (1 - wp) / 2
        upper_qp = 1 - lower_qp
        lower_quant = quantile(cdf, quantbound, support; p=lower_qp)[1]
        @warn "Can't compute 95% credible intervel, do $(Int(wp*100))% instead."
    end
    upper_quant = quantile(cdf, quantbound, support; p=upper_qp)[1]
    err = abs(cdf(upper_quant) -  cdf(lower_quant) - wp)/wp
    return [lower_quant, upper_quant], wp, err
end

function credible_interval(cdf::Function, quantbound::Function, support::Array{T,1}, ::Val{:narrow};
                            pdf=nothing, pdf_deriv=nothing, wp::T=.95) where T<:Float64
  #=
  Brief idea: bisection
    Suppose the target interval is [alpha*, beta*], i.e. integral of pdf on [alpha*, beta*] = wp
    and the corresponding height is pdf(alpha*) = pdf(beta*) = l_height*
    Say l* is the horizontal line intersecting with pdf curve at alpha* and beta*,
    we first find two horizontal l lines sitting on each side of l*
    One is lower with height l_height_low, and one is higher with height l_height_high
    Then we use a bisection way to gradually find the target l*
  =#

  #= Notations:
  [alpha, beta]:
      interval for integration, target credible interval
      s.t. 1) pdf(alpha) = pdf(beta) and 2) integral of pdf from alpha to beta = wp (like 0.95)
  bound:
      the bound of support of the pdf, currently assume this is available
      i.e. pdf(x) = 0 if z > bound
  l_height_low/l_height_high:
      height of the l line that is lower/higher than the target one
  alpha_low/alpha_high, beta_low/beta_high:
      the corresponding interval
  int_low/int_high:
      the corresponding integral value, i.e. integral of pdf on [alpha_low/high, beta_low/high]
  hquadrature:
      function from the package Cubature for numerical integration
      e.g. integral_heightue, error = hquadrature(f, a, b) gives integral of f on [a,b] with some error
      we mostly use this to compute integral of pdf on [alpha, beta] for different alpha and beta
  =#

 # helper functions
  #=
  given a height l_height, find the two intersections s.t. pdf(α) = pdf(β) = l_height
  Input:
    l_height: given height
    α_intvl/β_intvl: interval for root finding for alpha/beta
  Output:
    [α, β]: the two intersections
    int: integral value of pdf on [α, β]
  =#
  function find_αβ(l_height, α_intvl, β_intvl)
    # find α and β within ginve intervals α_intvl and β_intvl respectively
    temp_fun(x) = pdf(x) - l_height
    α = fzero(temp_fun,  α_intvl[1], α_intvl[2])
    β = fzero(temp_fun,  β_intvl[1], β_intvl[2])
    int = hquadrature(x -> pdf(x), α, β)[1]
    return α, β, int
  end

  #=
  Adjust the height in case initial choice of low/high is not proper
    Since we want int_low > wp and int_high < wp to do bisection,
    we have to adjust height if initial int_low < wp or int_high > wp
  Input:
    l: current height
    [α, β]: current interval
    int: current integral value
    MODE: indicates we are adjusting the low line or the high line
  Output:
    According values after adjustment
    s.t. int > wp if MODE == "low"; int < wp else.
 =#
  function adjust(l, α, β, int, MODE)
    # adjust if int < wp
    if MODE == "low"
        while int < wp
            l /= 2
            α, β, int = find_αβ(l, [support[1], α], [β, support[2]])
        end
    else
        while int > wp
            l = (l + pdf(mode_d))/2
            α, β, int = find_αβ(l, [α, mode_d], [mode_d, β])
        end
    end
    return l, α, β, int
  end
  mode_d = mode(pdf, cdf, support)
  l_height_low = pdf(support[1])
  α_low = support[1]
  β_low = fzero(x -> pdf(x) - l_height_low,  mode_d, support[2])
  int_low = hquadrature(x -> pdf(x), α_low, β_low)[1]
  # α_low, β_low, int_low = find_αβ(l_height_low, [support[1], mode_d], [mode_d, support[2]])
  # adjust height if l low is not lower than l* (i.e. int < wp)
  # l_height_low, α_low, β_low, int_low = adjust(l_height_low, α_low, β_low, int_low, "low")
  l_height_high = pdf(mode_d)*0.9
  α_high, β_high, int_high = find_αβ(l_height_high, [α_low, mode_d], [mode_d, β_low])
  # adjust height if l_high is not higher than l* (i.e. int > wp)
  l_height_high, α_high, β_high, int_high = adjust(l_height_high, α_high, β_high, int_high, "high")

  α_mid = 0.
  β_mid = 0.
  l_mid = 0.
  N = 0
  int = int_low
  while !isapprox(int, wp) && N < 50
  l_mid = (l_height_high + l_height_low)/2
  α_mid, β_mid, int = find_αβ(l_mid, [α_low, α_high], [β_high, β_low])
  if int > wp
      l_height_low = l_mid
      α_low = α_mid
      β_low = β_mid
  else
      l_height_high = l_mid
      α_high = α_mid
      β_high = β_mid
  end
  N += 1
  end
  err1 = abs(cdf(β_mid) -  cdf(α_mid) - wp)/wp
  err = max(err1, abs(pdf(α_mid)- pdf(β_mid))/abs(pdf(β_mid)))
  return ([α_mid, β_mid], err)

end

# @doc raw"""
# """
# function map_estimate(btg::BTG)
#     # TODO
# end

# @doc raw"""
# """
# function cross_validate(btg::BTG)
#     # TODO
# end


SNN = SpikingNeuralNetworks
param = SNN.IZParameter_more(;a = 0.02, b = 0.2, c = -65.0, d = 8.0, C=100.0, vr=-65.0, k=4.0, vt=-50.0, vPeak=20.0)


####### SIMULATION OF MARGINAL ###########################################

### +++ INITIALIZATION +++

### import packages
using Random, LinearAlgebra, RDatasets, MixedModels, Distributions, Optim, CSV, Dates, Statistics, CategoricalArrays

### define functions
vcov2 = function(x, xbar, σ, τ, n, γ)
    m = length(xbar)
    a = γ / τ^2 * m # 1-1 entry
    d = (sum(x.^2) - n * γ * sum(xbar.^2) ) / σ^2 # 2-2 entry
    b = γ / τ^2 / n * sum(x) # 2-1 entry
    [d -b; -b a] / (a * d - b^2) # return inverse (X^t V^-1 X)^-1
end

beta = function(vcov1, x, xbar, y, ybar, σ, τ, n, γ)
    c1 = γ / τ^2 * sum(ybar) # = 1^t V^-1 y
    c2 = (x' * y - n * γ * sum(ybar .* xbar))/σ^2 # = x^t V^-1 y
    vcov1 * [c1 c2]'  # (X^tV^-1X)^-1 X^tV^-1y
end

Q = function(data, y, ybar, x, xbar, n, N, μ0, quantile)

    ### fit model
    model = fit(MixedModel, @formula(y ~ x + (1 | v)), data)

    ### make sure the variance parameters are positive
    τ = maximum([VarCorr(model).σρ.v.σ[1], 0.001])
    σ = maximum([VarCorr(model).s, 0.001])

    ###     MARGINAL

    ### helpful parameters
    γ = τ^2 / (τ^2 + σ^2 / n)
    vcov1 = vcov2(x, xbar, σ, τ, n, γ) # = (X^tV^-1X)^-1
    d = (1-γ) .* hcat(fill(1, m), xbar) # d' = l' - b'X
    dd = d' * d
    a = n * τ^2 + σ^2

    ### assemble diagonal matrix parts
    g1 = γ * σ^2 / n
    g3 = 2 * σ^4 / τ^2 * γ / N/ (n - 1)

    ### assemble inverse of covariance matrix with Woodbury
    Σdiag = g1 + 2 * g3
    #Σ = Σdiag * I + d * vcov1 * d'
    Σi = (I - d * inv(Σdiag * inv(vcov1) + dd) * d' ) / Σdiag

    ###     CONDITIONAL

    ### assemble diagonal matrix parts
    l1 = γ^2 / n * σ^2

    l3_1 = 2 * σ^2 * τ^2 / N/ a^3
    l3_2 = 8 / N/ (n - 1) / a^4 / n^2 * τ^2 * σ^2 * (
        + (a^2 - σ^2) / n * (
            + σ^6 / a^2 * n
            - (σ^4  / a^2) * (τ + a / 2 * n)
            - (τ^2 - a / 2 + σ^2 * (1 + n^2)) * (n - 1))
        + σ^10 / a^2 * (n^3 - (n - 1)^2 + 1 / n)
        + a * σ^4 * (n - 1 + 1 / n / 2)  +
        - σ^6 * (2 * n - 1 + 1 / n) +
        + σ^8 / a * (1 + (n-1) * (n^2 + 1) / n) / 2)
    l3_3 = - 8 / N/ n * τ^2 * σ^4 / a^4 * (n - 1) * (a^2 - σ^4)
    l3_4 = - 8 / N/ n * τ^2 * σ^4 / a^3 / (n - 1) * (
        + a * ((n - 1) * n - 1)
        + σ^2 * (n + 1))
    l3 = 2 * (l3_1 + l3_2 + l3_3 + l3_4)

    l4 = 2 * σ^4 / N / a^2 * ( 1 / n + ((n - 1) * τ^2 +
            (σ^2 + τ^2)^2) / ((n - 1) * a^2))

    l5 = 2 * n / (n - 1) / N * σ^4 * (σ^2 - 2 * n * τ^2) / a^2

    ### assemble inverse of covariance matrix with Woodbury
    cΣdiag = maximum([l1 + l3 + l4 - l5, 0.001])
    cΣi = (I - d * inv(cΣdiag * σ^2 *
                inv(vcov1 * (2 * γ * I + n^2 * dd * vcov1)) + dd) * d' ) / cΣdiag

    ###     COMPUTE MEAN
    β = beta(vcov1, x, xbar, y, ybar, σ, τ, n, γ) # = fixef(model)
    μ = (1 - γ) * (β[1] .+ β[2] .* xbar) + γ * ybar

    ###     SAVE PIVOTS
    MarginalPivot = sum((μ - μ0)' .* ((μ - μ0)' * Σi), dims = 2)[1]
    ConditionalPivot = sum((μ - μ0)' .* ((μ - μ0)' * cΣi), dims = 2)[1]

    if quantile == true

        ###     NONCENTRALITY PARAMETER
        ay = d * β - (1 - γ) * ybar
        λ = maximum([(ay' * cΣi * ay)[1] - tr(cΣi) * σ^2 / a^2 / n, 0])

        ###     QUANTILES
        qc = cquantile(NoncentralChisq(m, λ), 0.05) # quantile of alpha = 0.05

        return [MarginalPivot ConditionalPivot qc β[1] β[2] τ σ Σdiag cΣdiag]
    else
        return [MarginalPivot ConditionalPivot]
    end
end

### fix parameters
p = 2
nsim = 5000
R = 300
Random.seed!(2021)

### +++ PER SPECIFICATION SETTING +++

### to compile
τ0 = 8^0.5
σ0 = 4^0.5
m = 50
n = 10
REdistr = "normal"
Edistr = "normal"

### variable parameters
N = m * n

β0 = [-4.9, 0.03]
x = 10 * randn(N)
xbar = vec(sum(reshape(x, n, :), dims=1) / n)

qm = cquantile(Chisq(m), 0.05) # quantile of alpha = 0.05

### conditional evaluation
if REdistr == "normal"
    v = τ0 * randn(m)
elseif REdistr == "tdist"
    v = rand(TDist(2), m)
elseif REdistr == "chi2"
    v = rand(Chisq(τ0^2 / 2), m) .- τ0^2 / 2
end
u = β0[1] .+ β0[2] * x + repeat(v, inner = n)
μ0 = β0[1] .+ β0[2] .* xbar + v

### obtain conditional variance
σ1 = zeros(10000, 1)
τ1 = zeros(10000, 1)
for i in 1:10000
    if Edistr == "normal"
        e = σ0 * randn(N)
    elseif Edistr == "tdist"
        e = rand(TDist(2), N)
    elseif Edistr == "chi2"
        e = rand(Chisq(σ0^2 / 2), N) .- σ0^2 / 2
    end
    y = u + e # conditional
    data = DataFrame(x = x, v = repeat(categorical(v), inner = n), y = y)
    model = fit(MixedModel, @formula(y ~ x + (1 | v)), data)
    τ1[i] = VarCorr(model).σρ.v.σ[1]
    σ1[i] = VarCorr(model).s
end
σ1 = mean(σ1) # should be (close to) σ0, but theory for conditional convergence is missing
τ1 = mean(τ1) # should be (close to) τ0, but theory for conditional convergence is missing

### simulation
Coverage = zeros(6)
λ = zeros(nsim)
Volume = ones(nsim, 6)
QM1 = zeros(nsim)
QC1 = zeros(nsim)
QM2 = zeros(nsim)
QC2 = zeros(nsim)
for i in 1:nsim

    ### +++ EMPIRICAL +++

    ### variable data
    if Edistr == "normal"
        e = σ0 * randn(N)
    elseif Edistr == "tdist"
        e = rand(TDist(2), N)
    elseif Edistr == "chi2"
        e = rand(Chisq(σ0^2 / 2), N) .- σ0^2 / 2
    end
    y = u + e # conditional
    ybar = vec(sum(reshape(y, n, :), dims=1) / n)
    data = DataFrame(x = x, v = repeat(categorical(v), inner = n), y = y)

    MP_bt_Alg1 = zeros(R-1)
    CP_bt_Alg1 = zeros(R-1)

    β1 = 1;  β2 = 1; τ = 1;  σ = 1

    MP, CP, qc, β1, β2, τ, σ, Σdiag, cΣdiag= Q(data, y, ybar, x, xbar, n, N, μ0, true)

    for r in 2:R

        v_bt = τ * randn(m)
        u_bt = β1 .+ β2 * x + repeat(v_bt, inner = n)
        μ0_bt = β1 .+ β2 .* xbar + v_bt

            ### generate new bootstrap data
            y_bt = u_bt + σ * randn(N) # conditional
            ybar_bt = vec(sum(reshape(y_bt, n, :), dims = 1) / n)
            data_bt = DataFrame(x = x, v = repeat(categorical(v_bt), inner = n), y = y_bt)

            MP_bt_Alg1[r-1], CP_bt_Alg1[r-1] = Q(data_bt, y_bt, ybar_bt, x, xbar, n, N, μ0_bt, false)
    end

    MP_bt_Alg2 = zeros(R-1)
    CP_bt_Alg2 = zeros(R-1)

    v_bt = τ * randn(m)
    u_bt = β1 .+ β2 * x + repeat(v_bt, inner = n)
    μ0_bt = β1 .+ β2 .* xbar + v_bt

    for r in 2:R

            ### generate new bootstrap data
            y_bt = u_bt + σ * randn(N) # conditional
            ybar_bt = vec(sum(reshape(y_bt, n, :), dims = 1) / n)
            data_bt = DataFrame(x = x, v = repeat(categorical(v_bt), inner = n), y = y_bt)

            MP_bt_Alg2[r-1], CP_bt_Alg2[r-1] = Q(data_bt, y_bt, ybar_bt, x, xbar, n, N, μ0_bt, false)
    end

    ###     COMPUTE BOOTSTRAP QUANTILES
    QM1[i] = quantile!(MP_bt_Alg1, 0.95)
    QC1[i] = quantile!(CP_bt_Alg1, 0.95)

    ###     COMPUTE BOOTSTRAP QUANTILES
    QM2[i] = quantile!(MP_bt_Alg2, 0.95)
    QC2[i] = quantile!(CP_bt_Alg2, 0.95)

    ###     COMPUTE COVERAGE
    Coverage[1] += (MP .< qm) / nsim
    Coverage[2] += (CP .< qc) / nsim
    Coverage[3] += (MP .< QM1[i]) / nsim
    Coverage[4] += (CP .< QC1[i]) / nsim
    Coverage[5] += (MP .< QM2[i]) / nsim
    Coverage[6] += (CP .< QC2[i]) / nsim

    ###     COMPUTE RELATIVE VOLUME
    # (via approximations as the determinant of large matrices is numerically unstable)
    Volume[i, 2] = (Σdiag / cΣdiag * qm / qc)^(-m / 2)
    Volume[i, 3] = (qm / QM1[i])^(-m / 2)
    Volume[i, 4] = (Σdiag / cΣdiag * qc / QC1[i])^(-m / 2)
    Volume[i, 5] = (qm / QM2[i])^(-m / 2)
    Volume[i, 6] = (Σdiag / cΣdiag * qc / QC2[i])^(-m / 2)

end


import Catalyst
import Catalyst:@reaction_network
import ModelingToolkit
using Transducers

function gene_expression_system(; mean_s=50, corr_time_s=1.0, corr_time_x=0.1, u0=SA[mean_s, mean_s], dtimes=0:0.1:2.0)
    sn = @reaction_network begin
        κ, ∅ --> S
        λ, S --> ∅
    end κ λ

    xn = @reaction_network begin
        ρ, S --> X + S
        μ, X --> ∅
    end ρ μ

    λ = 1 / corr_time_s
    κ = mean_s * λ
    μ = 1 / corr_time_x
    ρ = μ

    ps = [κ, λ]
    px = [ρ, μ]

    SimpleSystem(sn, xn, u0, ps, px, dtimes)
end

function chemotaxis_system(;
        mean_L=20,
        num_receptors=10000,
        Y_tot=5000,
        L_timescale=1.0,
        LR_timescale=0.01,
        LR_ratio=0.5,
        Y_timescale=0.1,
        Y_ratio=1/6,
        q=0,
        dtimes=0:0.1:2.0
    )
    mean_LR = num_receptors * LR_ratio
    mean_R = num_receptors - mean_LR

    mean_Yp = Y_tot * Y_ratio
    mean_Y = Y_tot - mean_Yp

    eq_L = mean_L * exp(-q)

    sn = @reaction_network begin
        κ, ∅ --> L
        λ, L --> ∅
    end κ λ

    rn = @reaction_network begin
        ρ, L + R --> L + LR
        μ, LR --> R
    end ρ μ

    xn = @reaction_network begin
        δ, LR + Y --> Yp + LR
        χ, Yp --> Y
    end δ χ

    u0 = round.([mean_L, mean_R, mean_LR, mean_Y, mean_Yp])
    ps = [mean_L, 1 / L_timescale]

    ρ = inv(eq_L * LR_timescale * (1 + mean_R / mean_LR))
    μ = eq_L * ρ * mean_R / mean_LR

    pr = [ρ, μ]

    χ = inv(Y_timescale * (1 + Y_ratio))
    δ = χ * Y_ratio / mean_LR

    px = [δ, χ]

    ComplexSystem(sn, rn, xn, u0, ps, pr, px, dtimes)
end

function chemotaxis_parameters(;
    E₀ = 3.0,
    Kₐ = 500,
    Kᵢ = 25,
    δg = log(Kₐ/Kᵢ),
    δf = -1.5,
    k⁺ = 0.05,
    k⁺ₐ = k⁺,
    k⁺ᵢ = k⁺,
    k⁻ₐ = Kₐ * k⁺ₐ,
    k⁻ᵢ = Kᵢ * k⁺ᵢ,
    a_star = 0.5,
    γ = 1 / 10,
    k_B = (1 - a_star) * γ / abs(δf),
    k_R =      a_star  * γ / abs(δf)
)
    [
        E₀,
        δg,
        δf,
        k⁺ₐ,
        k⁺ᵢ,
        k⁻ₐ,
        k⁻ᵢ,
        k_B,
        k_R
    ]
end

function cooperative_chemotaxis_system(;
    lmax = 3,
    mmax = 9,
    n_clusters = 100,
    n_chey = 10_000,

    mean_l = 50,
    tau_l = 1.0,

    phosphorylate = 2000.0 / (n_chey * n_clusters),
    dephosphorylate = 2000.0 / (n_chey),
    dtimes = 0:0.1:20.0,
    varargs...
)
    sn = @reaction_network begin
        κ, ∅ --> L
        λ, L --> ∅
    end κ λ

    rn = Catalyst.make_empty_network()

    @Catalyst.parameters t E0 δg δf lba lbi lda ldi mda mbi ρ
    @Catalyst.variables L(t) Y(t) Yp(t)

    Catalyst.addspecies!(rn, L)

    Catalyst.addparam!(rn, E0)
    Catalyst.addparam!(rn, δg)
    Catalyst.addparam!(rn, δf)
    Catalyst.addparam!(rn, lba)
    Catalyst.addparam!(rn, lbi)
    Catalyst.addparam!(rn, lda)
    Catalyst.addparam!(rn, ldi)
    Catalyst.addparam!(rn, mda)
    Catalyst.addparam!(rn, mbi)

    xn = @reaction_network begin
        μ, Yp --> Y
    end μ

    Catalyst.addparam!(xn, ρ)

    spmap = Dict()
    for l=0:lmax, m=0:mmax
        receptor_species = ModelingToolkit.Num(ModelingToolkit.Variable{ModelingToolkit.FnType{Tuple{Any},Real}}(Symbol("R_", l, "_", m)))(t)

        spmap[(l, m)] = receptor_species

        Catalyst.addspecies!(xn, receptor_species)
        Catalyst.addspecies!(rn, receptor_species)
    end

    p_active(l, m) = 1 / (1 + exp(E0 + l*δg + m*δf))

    for l=0:lmax, m=0:mmax

        if l > 0
            ligand_bind = Catalyst.Reaction((lba * p_active(l-1, m) + lbi * (1 - p_active(l-1, m))) * (lmax + 1 - l), [spmap[(l-1, m)], L], [spmap[(l, m)], L])
            ligand_unbind = Catalyst.Reaction((lda * p_active(l, m) + ldi * (1 - p_active(l, m))) * l, [spmap[(l, m)]], [spmap[(l-1, m)]])
            Catalyst.addreaction!(rn, ligand_bind)
            Catalyst.addreaction!(rn, ligand_unbind)
        end

        if m > 0
            demethylate_active = Catalyst.Reaction(mda * p_active(l, m), [spmap[(l, m)]], [spmap[(l, m-1)]])
            methylate_inactive = Catalyst.Reaction(mbi * (1 - p_active(l, m-1)), [spmap[(l, m-1)]], [spmap[(l, m)]])

            Catalyst.addreaction!(rn, demethylate_active)
            Catalyst.addreaction!(rn, methylate_inactive)
        end

        # every receptor phosphorylates Y with rate ρ if the receptor is active
        # if the receptor is inactive, no phosphorylation can happen
        receptor = spmap[(l, m)]
        phosphorylation = Catalyst.Reaction(ρ * p_active(l, m), [Y, receptor], [Yp, receptor])
        Catalyst.addreaction!(xn, phosphorylation)
    end

    joint = merge(merge(sn, rn), xn)

    u0 = zeros(Int, length(spmap) + 3)
    u0[1] = round(Int, mean_l)
    u0[Catalyst.speciesmap(joint)[spmap[(0, 0)]]] = n_clusters
    u0[Catalyst.speciesmap(joint)[Y]] = n_chey

    ps = [mean_l/tau_l, 1.0/tau_l]
    pr = chemotaxis_parameters(; varargs...)
    px = [dephosphorylate, phosphorylate]

    ComplexSystem(sn, rn, xn, u0, ps, pr, px, dtimes; aggregator=DiffEqJump.RSSACR())
end

getname(sym) = String(ModelingToolkit.operation(sym).name)

function parse_receptor(species)
    sname = getname(species)
    mtch = match(r"([A-Z])_(\d+)_(\d+)", sname)
    if mtch !== nothing
        l = parse(Int, mtch.captures[2])
        m = parse(Int, mtch.captures[3])
        (l, m)
    else
        nothing
    end
end

function receptor_states(rs::ReactionSystem)
    smap = Catalyst.speciesmap(rs)
    xf = KeepSomething() do (species, index)
        result = parse_receptor(species)
        if result === nothing
            nothing
        else
            result => index
        end
    end
    smap |> xf
end

receptor_states(system::ComplexSystem) = receptor_states(reaction_network(system))

function active_receptors(conf::SRXconfiguration, system::ComplexSystem)
    rstates = receptor_states(system)

    E0 = system.pr[1]
    δg = system.pr[2]
    δf = system.pr[3]

    p_active(l, m) = 1 / (1 + exp(E0 + l*δg + m*δf))

    p_a = rstates |> Map(((l, m), i)::Pair -> i => p_active(l, m)) |> collect

    f = function(u)
        sum(p_a) do (i, p)
            p * u[i]
        end
    end

    merge_trajectories(conf.s_traj, conf.r_traj, conf.x_traj) |> Map((u,t,i)::Tuple -> (ensurevec(f(u)),t,i)) |> collect_trajectory
end

precompile(gene_expression_system, ())
precompile(chemotaxis_system, ())
precompile(cooperative_chemotaxis_system, ())


function FFTOp(in,adj;normalize=true)
	norm = normalize ? sqrt(length(in[:])) : 1.
	if (adj)
		out = fft(in)/norm
	else
		out = bfft(in)/norm
	end

	return out
end


using CompScienceMeshes
using BEAST

Γ1 = readmesh(joinpath(@__DIR__,"sphere2.in"))
Γ2 = CompScienceMeshes.SComplex2D(Γ1)
# Y = raviartthomas(Γ1)
Nd1 = BEAST.nedelec(Γ1)
Nd2 = BEAST.nedelec(Γ2)
RT1 = n × Nd1
RT2 = n × Nd2

N1 = skeleton(Γ1,0)
N2 = skeleton(Γ2,0)

for (n1,n2) in zip(N1.faces, N2.nodes) @assert n1[1] == n2[1] end

E1 = skeleton(Γ1,1)
E2 = skeleton(Γ2,1)

using LinearAlgebra
for (e1,e2) in zip(E1.faces, E2.edges)
    q = getindex.(Ref(E2.nodes), e2)
    @assert q[1][1] == e1[1]
    @assert q[2][1] == e1[2]
    ctr1 = cartesian(center(chart(E1, e1)))
    ctr2 = cartesian(center(chart(E2, e2)))
    if norm(ctr1-ctr2) > 1e-8
        @show ctr1
        @show ctr2
        error("stop")
    end
end

κ = 1.0
t = Maxwell3D.singlelayer(wavenumber=κ)
E = Maxwell3D.planewave(direction=ẑ, polarization=x̂, wavenumber=κ)
e = (n × E) × n

Z1 = assemble(t,RT1,RT1)
Z2 = assemble(t,RT2,RT2)

@hilbertspace j
@hilbertspace k
efie = @discretise t[k,j]==e[k]  j∈RT2 k∈RT2
u = gmres(efie)

Γ = Γ2
X = RT2

include("utils/postproc.jl")
include("utils/plotresults.jl")


module AbstractSphericalHarmonics

using ComputedFieldTypes
using LinearAlgebra
using StaticArrays

################################################################################

# FastSphericalHarmonics requires julia >= 1.7
# SSHT requires julia >= 1.6

# using FastSphericalHarmonics
using SSHT

export ash_grid_size, ash_nmodes
export ash_ntheta, ash_nphi, ash_thetas, ash_phis, ash_point_coord, ash_point_delta, ash_grid_as_phi_theta
export ash_mode_index, ash_mode_numbers
export ash_transform!, ash_transform, ash_evaluate!, ash_evaluate
export ash_eth!, ash_eth, ash_ethbar!, ash_ethbar

################################################################################

bitsign(i::Integer) = isodd(i) ? -one(i) : one(i)

################################################################################

include("sylm.jl")
include("tensor.jl")

end


module C1_1
function f(q::Vector{T}, q_dot::Vector{T}, xi::T) where T
    l1 = q[1]
    l2 = q[2]
    l3 = q[3]
    l1_dot = q_dot[1]
    l2_dot = q_dot[2]
    l3_dot = q_dot[3]

    l1_dot.*(7.00120879963642e-11*l1 - 4.28998384220961e-11*l2 - 2.75081888864459e-11*l3 + 6.01947273729361e-14) - l2_dot.*(4.28998384220961e-11*l1 - 5.44544999423272e-11*l2 + 1.12667056568291e-11*l3 + 4.41800310108094e-14) - l3_dot.*(6.17485674830662e-12*l1 + 4.31447849383884e-11*l2 - 4.92763081818322e-11*l3 - 6.44040569580091e-15)
end
end


using SpecialFunctions: erfcx
import RationalApproximations: locate_extrema_idx, RationalFunc, aaa_v2, minimax
using PyPlot

G(t) = (1-t)/(1+t)
α = 0.5
β = 1.0
fα(t) = erfcx(G(t))

t = range(-1, 1, length=800)
max_m = 5
tol = 1e-10
F = fα.(t)
r, err = aaa_v2(F, t, max_m, tol)
m = length(r)
resid = F - r[m].(t)
idx = locate_extrema_idx(resid, m)

figure(1)
plot(t, resid, t[idx], resid[idx], "o")
grid(true)
xlabel(L"$t$", fontsize=12)
xylims=axis()
savefig("figure13L.pdf")

supp_pt = r[m].supp_pt
max_iterations=3
tx = t[idx]
opt_r, zmin, zmax = minimax(fα, tx, supp_pt, (-1.0, 1.0), max_iterations)
resid = F - opt_r[end].(t)
figure(2)
plot(t, resid, tx, fα.(tx) - opt_r[end].(tx), "o")
grid(true)
xlabel(L"$t$", fontsize=12)
axis(xylims)
savefig("figure13R.pdf")


module PolymakeOscar

using Polymake, Hecke

Hecke.nrows(A::Polymake.pm_MatrixAllocated) = Int(size(A)[1])
Hecke.ncols(A::Polymake.pm_MatrixAllocated) = Int(size(A)[2])

#solves Ax <= b
function solve_ineq(A::fmpz_mat, b::fmpz_mat)
  bA = Array{BigInt, 2}(hcat(-b, A))
  p = @pm Polytope.Polytope( :INEQUALITIES => bA)
  inner = p.INTERIOR_LATTICE_POINTS
  out = p.BOUNDARY_LATTICE_POINTS

  res = zero_matrix(FlintZZ, nrows(inner) + nrows(out), ncols(A))
  for i=1:nrows(out)
    @assert out[i,1] == 1
    for j=1:ncols(A)
      res[i,j] = out[i, j+1]
    end
  end
  for i=1:nrows(inner)
    @assert inner[i,1] == 1
    for j=1:ncols(A)
      res[i+nrows(out), j] = inner[i, j+1]
    end
  end
  return res
end

function solve_non_negative(A::fmpz_mat, b::fmpz_mat)
  bA = Array{BigInt, 2}(hcat(-b, A))
  zI = Array{BigInt, 2}(hcat(zero_matrix(FlintZZ, ncols(A), 1), MatrixSpace(FlintZZ, ncols(A), ncols(A))(1)))
  p = @pm Polytope.Polytope(:EQUATIONS => deepcopy(bA), :INEQUALITIES => deepcopy(zI))
  inner = p.INTERIOR_LATTICE_POINTS
  out = p.BOUNDARY_LATTICE_POINTS

  res = zero_matrix(FlintZZ, nrows(inner) + nrows(out), ncols(A))
  for i=1:nrows(out)
    @assert out[i,1] == 1
    for j=1:ncols(A)
      res[i,j] = out[i, j+1]
    end
  end
  for i=1:nrows(inner)
    @assert inner[i,1] == 1
    for j=1:ncols(A)
      res[i+nrows(out), j] = inner[i, j+1]
    end
  end
  return res
end

function solve_mixed(A::fmpz_mat, b::fmpz_mat, C::fmpz_mat)  # Ax == b && Cx >= 0
  bA = Array{BigInt, 2}(hcat(-b, A))
  z = findall(i->!iszero_row(C, i), 1:nrows(C))
  zI = Array{BigInt, 2}(hcat(zero_matrix(FlintZZ, nrows(C), 1), C))[z, :]
  p = @pm Polytope.Polytope(:EQUATIONS => bA, :INEQUALITIES => zI)
  inner = p.INTERIOR_LATTICE_POINTS
  out = p.BOUNDARY_LATTICE_POINTS

  res = zero_matrix(FlintZZ, nrows(inner) + nrows(out), ncols(A))
  for i=1:nrows(out)
    @assert out[i,1] == 1
    for j=1:ncols(A)
      res[i,j] = out[i, j+1]
    end
  end
  for i=1:nrows(inner)
    @assert inner[i,1] == 1
    for j=1:ncols(A)
      res[i+nrows(out), j] = inner[i, j+1]
    end
  end
  return res
end

function Hecke.valuation(a::FacElem{fmpz, FlintIntegerRing}, p::fmpz)
  return sum(k*valuation(b, p) for (b, k) = a.fac)
end

function norm_equation2_fac_elem(R::NfAbsOrd, k::fmpz; abs::Bool = false)
  @assert Hecke.ismaximal(R)
  lp = factor(k*R)
  s, ms = Hecke.sunit_mod_units_group_fac_elem(collect(keys(lp)))
  C = vcat([matrix(FlintZZ, 1, ngens(s), [valuation(ms(s[i]), p) for i=1:ngens(s)]) for p = keys(lp)])

  lp = factor(k)
  A = vcat([matrix(FlintZZ, 1, ngens(s), [valuation(Hecke.factored_norm(ms(s[i])), p) for i=1:ngens(s)]) for p = keys(lp.fac)])
  b = matrix(FlintZZ, length(lp.fac), 1, [valuation(k, p) for p = keys(lp.fac)])

  so = solve_mixed(A, b, C)
  sol = [ms(s(sub(so, i:i, 1:ncols(so)))) for i=1:nrows(so)]

  if !abs
    u, mu = unit_group_fac_elem(R)
    i = findfirst(x -> norm(mu(x)) == -1, gens(u))
    ns = [norm(x) for x = sol]
    if i === nothing
      return [sol[i] for i in 1:length(sol) if ns[i] == k]
    end
    U = mu(u[i])
    return [ ns[i] == k ? sol[i] : sol[i] * U for i = 1:length(sol)]
  end
  return sol
end

function norm_equation_fac_elem(R::NfAbsOrd, k::fmpz; abs::Bool = false)
  @assert Hecke.ismaximal(R)
  lp = factor(k)
  S = []
  for (p, k) = lp.fac
    P = prime_decomposition(R, p)
    s = solve_non_negative(matrix(FlintZZ, 1, length(P), [degree(x[1]) for x = P]), matrix(FlintZZ, 1, 1, [k]))
    push!(S, (P, [view(s, i:i, 1:ncols(s)) for i=1:nrows(s)]))
  end
  sol = []
  for x in Base.Iterators.ProductIterator(Tuple(t[2] for t = S))
    I = ideal(R, 1)
    for i = 1:length(S)
      I *= prod(S[i][1][j][1]^Int(x[i][j]) for j=1:length(S[i][1]))
    end
    fl, g = Hecke.isprincipal_fac_elem(I)
    if fl
      push!(sol, g)
    end
  end
  if !abs
    u, mu = unit_group_fac_elem(R)
    i = findfirst(x -> norm(mu(x)) == -1, gens(u))
    ns = [norm(x) for x = sol]
    if i === nothing
      return [sol[i] for i in 1:length(sol) if ns[i] == k]
    end
    U = mu(u[i])
    return [ ns[i] == k ? sol[i] : sol[i] * U for i = 1:length(sol)]
  end
  return sol
end

norm_equation_fac_elem(R::NfAbsOrd, k::Integer; abs::Bool = false) =
                            norm_equation_fac_elem(R, fmpz(k), abs = abs)

function norm_equation(R::NfAbsOrd, k::fmpz; abs::Bool = false)
  s = norm_equation_fac_elem(R, k, abs = abs)
  return [R(evaluate(x)) for x = s]
end

norm_equation(R::NfAbsOrd, k::Integer; abs::Bool = false) = norm_equation(R, fmpz(k), abs = abs)

function norm_equation_fac_elem(R::Hecke.NfRelOrd{nf_elem,Hecke.NfOrdFracIdl}, a::NfAbsOrdElem{AnticNumberField,nf_elem})

  @assert Hecke.ismaximal(R)
  Ka, mKa, mkK = absolute_field(nf(R))
  Ra = maximal_order(Ka)
  class_group(Ra)
  k = nf(parent(a))
  class_group(parent(a))

  lp = factor(Ra(mkK(k(a)))*Ra)
  la = factor(a*parent(a))
  S, mS = Hecke.sunit_mod_units_group_fac_elem(collect(keys(lp)))
  s, ms = Hecke.sunit_mod_units_group_fac_elem(collect(keys(la)))
  No = hom(S, s, elem_type(s)[preimage(ms, norm(mkK, mS(g))) for g = gens(S)])
  q, mq = quo(S, kernel(No)[1])
  q, mms = snf(q)
  mq = mq*inv(mms)

  C = vcat([matrix(FlintZZ, 1, ngens(q), [valuation(mS(preimage(mq, q[i])), p) for i=1:ngens(q)]) for p = keys(lp)])

  A = vcat([matrix(FlintZZ, 1, ngens(q), [valuation(norm(mkK, mS(preimage(mq, g))), p) for g in gens(q)]) for p = keys(la)])
  b = matrix(FlintZZ, length(la), 1, [valuation(a, p) for p = keys(la)])

  so = solve_mixed(A, b, C)
  u, mu = Hecke.unit_group_fac_elem(parent(a))
  U, mU = Hecke.unit_group_fac_elem(Ra)
  No = hom(U, u, elem_type(u)[preimage(mu, norm(mkK, mU(g))) for g = gens(U)])
  sol = []
  for i = 1:nrows(so)
    aa = mS(preimage(mq, q(sub(so, i:i, 1:ncols(so)))))
    b = norm(mkK, aa)
    c = b*inv(FacElem(k(a)))
    d = preimage(mu, c)
    fl, p = haspreimage(No, d)
    if fl
      push!(sol, FacElem(Dict(mKa(x) => v for (x, v) = (aa*inv(mU(p))).fac)))
    end
  end

  return sol
end

end


using FactCheck, Base.Test, Images, Color, FixedPointNumbers

facts("Algorithms") do
    # Comparison of each element in arrays with a scalar
    approx_equal(ar, v) = all(abs(ar.-v) .< sqrt(eps(v)))
    approx_equal(ar::Images.AbstractImage, v) = approx_equal(Images.data(ar), v)

    context("Arithmetic") do
        img = convert(Images.Image, zeros(3,3))
        img2 = (img .+ 3)/2
        @fact all(img2 .== 1.5) --> true
        img3 = 2img2
        @fact all(img3 .== 3) --> true
        img3 = copy(img2)
        img3[img2 .< 4] = -1
        @fact all(img3 .== -1) --> true
        img = convert(Images.Image, rand(3,4))
        A = rand(3,4)
        img2 = img .* A
        @fact all(Images.data(img2) == Images.data(img).*A) --> true
        img2 = convert(Images.Image, A)
        img2 = img2 .- 0.5
        img3 = 2img .* img2
        img2 = img ./ A
        img2 = (2img).^2
        # Same operations with ColorValue images
        img = Images.colorim(zeros(Float32,3,4,5))
        img2 = (img .+ RGB{Float32}(1,1,1))/2
        @fact all(img2 .== RGB{Float32}(1,1,1)/2) --> true
        img3 = 2img2
        @fact all(img3 .== RGB{Float32}(1,1,1)) --> true
        A = fill(2, 4, 5)
        @fact all(A.*img2 .== fill(RGB{Float32}(1,1,1), 4, 5)) --> true
        img2 = img2 .- RGB{Float32}(1,1,1)/2
        A = rand(Uint8,3,4)
        img = reinterpret(Images.Gray{Ufixed8}, Images.grayim(A))
        imgm = mean(img)
        imgn = img/imgm
        @fact reinterpret(Float32, Images.data(imgn)) --> roughly(convert(Array{Float32}, A/mean(A)))
    end

    context("Reductions") do
        A = rand(5,5,3)
        img = Images.colorim(A, "RGB")
        s12 = sum(img, (1,2))
        @fact Images.colorspace(s12) --> "RGB"
        s3 = sum(img, (3,))
        @fact Images.colorspace(s3) --> "Unknown"
        A = [NaN, 1, 2, 3]
        @fact Images.meanfinite(A, 1) --> roughly([2])
        A = [NaN 1 2 3;
             NaN 6 5 4]
        @test_approx_eq Images.meanfinite(A, 1) [NaN 3.5 3.5 3.5]
        @test_approx_eq Images.meanfinite(A, 2) [2, 5]'
        @test_approx_eq Images.meanfinite(A, (1,2)) [3.5]
        @fact Images.minfinite(A) --> 1
        @fact Images.maxfinite(A) --> 6
        @fact Images.maxabsfinite(A) --> 6
        A = rand(Float32,3,5,5)
        img = Images.colorim(A, "RGB")
        dc = Images.data(Images.meanfinite(img, 1))-reinterpret(RGB{Float32}, mean(A, 2), (1,5))
        @fact maximum(map(abs, dc)) --> less_than(1e-6)

        a = convert(Array{Uint8}, [1, 1, 1])
        b = convert(Array{Uint8}, [134, 252, 4])
        @fact Images.sad(a, b) --> 387
        @fact Images.ssd(a, b) --> 80699
        af = reinterpret(Ufixed8, a)
        bf = reinterpret(Ufixed8, b)
        @fact Images.sad(af, bf) --> roughly(387f0/255)
        @fact Images.ssd(af, bf) --> roughly(80699f0/255^2)
        ac = reinterpret(RGB{Ufixed8}, a)
        bc = reinterpret(RGB{Ufixed8}, b)
        @fact Images.sad(ac, bc) --> roughly(387f0/255)
        @fact Images.ssd(ac, bc) --> roughly(80699f0/255^2)
        ag = reinterpret(RGB{Ufixed8}, a)
        bg = reinterpret(RGB{Ufixed8}, b)
        @fact Images.sad(ag, bg) --> roughly(387f0/255)
        @fact Images.ssd(ag, bg) --> roughly(80699f0/255^2)
    end

    context("fft and ifft") do
        A = rand(Float32, 3, 5, 6)
        img = Images.colorim(A)
        imgfft = fft(img)
        @fact Images.data(imgfft) --> roughly(fft(A, 2:3))
        @fact Images.colordim(imgfft) --> 1
        img2 = ifft(imgfft)
        @fact img2 --> roughly(reinterpret(Float32, img))
    end

    context("Array padding") do
        A = [1 2; 3 4]
        @fact Images.padindexes(A, 1, 0, 0, "replicate") --> [1,2]
        @fact Images.padindexes(A, 1, 1, 0, "replicate") --> [1,1,2]
        @fact Images.padindexes(A, 1, 0, 1, "replicate") --> [1,2,2]
        @fact Images.padarray(A, (0,0), (0,0), "replicate") --> A
        @fact Images.padarray(A, (1,2), (2,0), "replicate") --> [1 1 1 2; 1 1 1 2; 3 3 3 4; 3 3 3 4; 3 3 3 4]
        @fact Images.padindexes(A, 1, 1, 0, "circular") --> [2,1,2]
        @fact Images.padindexes(A, 1, 0, 1, "circular") --> [1,2,1]
        @fact Images.padarray(A, [2,1], [0,2], "circular") --> [2 1 2 1 2; 4 3 4 3 4; 2 1 2 1 2; 4 3 4 3 4]
        @fact Images.padindexes(A, 1, 1, 0, "symmetric") --> [1,1,2]
        @fact Images.padindexes(A, 1, 0, 1, "symmetric") --> [1,2,2]
        @fact Images.padarray(A, (1,2), (2,0), "symmetric") --> [2 1 1 2; 2 1 1 2; 4 3 3 4; 4 3 3 4; 2 1 1 2]
        @fact Images.padarray(A, (1,2), (2,0), "value", -1) --> [-1 -1 -1 -1; -1 -1 1 2; -1 -1 3 4; -1 -1 -1 -1; -1 -1 -1 -1]
        A = [1 2 3; 4 5 6]
        @fact Images.padindexes(A, 2, 1, 0, "reflect") --> [2,1,2,3]
        @fact Images.padindexes(A, 2, 0, 1, "reflect") --> [1,2,3,2]
        @fact Images.padarray(A, (1,2), (2,0), "reflect") --> [6 5 4 5 6; 3 2 1 2 3; 6 5 4 5 6; 3 2 1 2 3; 6 5 4 5 6]
        A = [1 2; 3 4]
        @fact Images.padarray(A, (1,1)) --> [1 1 2 2; 1 1 2 2; 3 3 4 4; 3 3 4 4]
        @fact Images.padarray(A, (1,1), "replicate", "both") --> [1 1 2 2; 1 1 2 2; 3 3 4 4; 3 3 4 4]
        @fact Images.padarray(A, (1,1), "circular", "pre") --> [4 3 4; 2 1 2; 4 3 4]
        @fact Images.padarray(A, (1,1), "symmetric", "post") --> [1 2 2; 3 4 4; 3 4 4]
        A = ["a" "b"; "c" "d"]
        @fact Images.padarray(A, (1,1)) --> ["a" "a" "b" "b"; "a" "a" "b" "b"; "c" "c" "d" "d"; "c" "c" "d" "d"]
        @fact_throws ErrorException Images.padindexes(A, 1, 1, 1, "unknown")
        @fact_throws ErrorException Images.padarray(A, (1,1), "unknown")
        # issue #292
        A = trues(3,3)
        @fact typeof(Images.padarray(A, (1,2), (2,1), "replicate")) --> BitArray{2}
        @fact typeof(Images.padarray(Images.grayim(A), (1,2), (2,1), "replicate")) --> BitArray{2}
    end

    context("Filtering") do
        EPS = 1e-14
        imgcol = Images.colorim(rand(3,5,6))
        imgcolf = convert(Images.Image{RGB{Ufixed8}}, imgcol)
        for T in (Float64, Int)
            A = zeros(T,3,3); A[2,2] = 1
            kern = rand(3,3)
            @fact maximum(abs(Images.imfilter(A, kern) - rot180(kern))) --> less_than(EPS)
            kern = rand(2,3)
            @fact maximum(abs(Images.imfilter(A, kern)[1:2,:] - rot180(kern))) --> less_than(EPS)
            kern = rand(3,2)
            @fact maximum(abs(Images.imfilter(A, kern)[:,1:2] - rot180(kern))) --> less_than(EPS)
        end
        kern = zeros(3,3); kern[2,2] = 1
        @fact maximum(map(abs, imgcol - Images.imfilter(imgcol, kern))) --> less_than(EPS)
        @fact maximum(map(abs, imgcolf - Images.imfilter(imgcolf, kern))) --> less_than(EPS)
        for T in (Float64, Int)
            # Separable kernels
            A = zeros(T,3,3); A[2,2] = 1
            kern = rand(3).*rand(3)'
            @fact maximum(abs(Images.imfilter(A, kern) - rot180(kern))) --> less_than(EPS)
            kern = rand(2).*rand(3)'
            @fact maximum(abs(Images.imfilter(A, kern)[1:2,:] - rot180(kern))) --> less_than(EPS)
            kern = rand(3).*rand(2)'
            @fact maximum(abs(Images.imfilter(A, kern)[:,1:2] - rot180(kern))) --> less_than(EPS)
        end
        A = zeros(3,3); A[2,2] = 1
        kern = rand(3,3)
        @fact maximum(abs(Images.imfilter_fft(A, kern) - rot180(kern))) --> less_than(EPS)
        kern = rand(2,3)
        @fact maximum(abs(Images.imfilter_fft(A, kern)[1:2,:] - rot180(kern))) --> less_than(EPS)
        kern = rand(3,2)
        @fact maximum(abs(Images.imfilter_fft(A, kern)[:,1:2] - rot180(kern))) --> less_than(EPS)
        kern = zeros(3,3); kern[2,2] = 1
        @fact maximum(map(abs, imgcol - Images.imfilter_fft(imgcol, kern))) --> less_than(EPS)
        @fact maximum(map(abs, imgcolf - Images.imfilter_fft(imgcolf, kern))) --> less_than(EPS)

        @fact approx_equal(Images.imfilter(ones(4,4), ones(3,3)), 9.0) --> true
        @fact approx_equal(Images.imfilter(ones(3,3), ones(3,3)), 9.0) --> true
        @fact approx_equal(Images.imfilter(ones(3,3), [1 1 1;1 0.0 1;1 1 1]), 8.0) --> true
        img = convert(Images.Image, ones(4,4))
        @fact approx_equal(Images.imfilter(img, ones(3,3)), 9.0) --> true
        A = zeros(5,5,3); A[3,3,[1,3]] = 1
        @fact Images.colordim(A) --> 3
        kern = rand(3,3)
        kernpad = zeros(5,5); kernpad[2:4,2:4] = kern
        Af = Images.imfilter(A, kern)

        @fact cat(3, rot180(kernpad), zeros(5,5), rot180(kernpad)) --> roughly(Af)
        Aimg = permutedims(convert(Images.Image, A), [3,1,2])
        @fact Images.imfilter(Aimg, kern) --> roughly(permutedims(Af, [3,1,2]))
        @fact approx_equal(Images.imfilter(ones(4,4),ones(1,3),"replicate"), 3.0) --> true

        A = zeros(5,5); A[3,3] = 1
        kern = rand(3,3)
        Af = Images.imfilter(A, kern, "inner")
        @fact Af --> rot180(kern)
        Afft = Images.imfilter_fft(A, kern, "inner")
        @fact Af --> roughly(Afft)
        h = [0.24,0.87]
        hfft = Images.imfilter_fft(eye(3), h, "inner")
        hfft[abs(hfft) .< 3eps()] = 0
        @fact Images.imfilter(eye(3), h, "inner") --> roughly(hfft)  # issue #204

        @fact approx_equal(Images.imfilter_gaussian(ones(4,4), [5,5]), 1.0) --> true
        A = fill(convert(Float32, NaN), 3, 3)
        A[1,1] = 1
        A[2,1] = 2
        A[3,1] = 3
        @fact Images.imfilter_gaussian(A, [0,0]) --> exactly(A)
        @test_approx_eq Images.imfilter_gaussian(A, [0,3]) A
        B = copy(A)
        B[isfinite(B)] = 2
        @test_approx_eq Images.imfilter_gaussian(A, [10^3,0]) B
        @fact maximum(map(abs, Images.imfilter_gaussian(imgcol, [10^3,10^3]) - mean(imgcol))) --> less_than(1e-4)
        @fact maximum(map(abs, Images.imfilter_gaussian(imgcolf, [10^3,10^3]) - mean(imgcolf))) --> less_than(1e-4)
        A = rand(4,5)
        img = reinterpret(Images.Gray{Float64}, Images.grayim(A))
        imgf = Images.imfilter_gaussian(img, [2,2])
        @fact reinterpret(Float64, Images.data(imgf)) --> roughly(Images.imfilter_gaussian(A, [2,2]))
        A = rand(3,4,5)
        img = Images.colorim(A)
        imgf = Images.imfilter_gaussian(img, [2,2])
        @fact reinterpret(Float64, Images.data(imgf)) --> roughly(Images.imfilter_gaussian(A, [0,2,2]))

        A = zeros(Int, 9, 9); A[5, 5] = 1
        @fact maximum(abs(Images.imfilter_LoG(A, [1,1]) - Images.imlog(1.0))) --> less_than(EPS)
        @fact maximum(Images.imfilter_LoG([0 0 0 0 1 0 0 0 0], [1,1]) - sum(Images.imlog(1.0),1)) --> less_than(EPS)
        @fact maximum(Images.imfilter_LoG([0 0 0 0 1 0 0 0 0]', [1,1]) - sum(Images.imlog(1.0),2)) --> less_than(EPS)

        @fact Images.imaverage() --> fill(1/9, 3, 3)
        @fact Images.imaverage([3,3]) --> fill(1/9, 3, 3)
        @fact_throws ErrorException Images.imaverage([5])
    end

    context("Restriction") do
        imgcol = Images.colorim(rand(3,5,6))
        A = reshape([convert(UInt16, i) for i = 1:60], 4, 5, 3)
        B = Images.restrict(A, (1,2))
        Btarget = cat(3, [ 0.96875  4.625   5.96875;
                           2.875   10.5    12.875;
                           1.90625  5.875   6.90625],
                      [ 8.46875  14.625 13.46875;
                        17.875    30.5   27.875;
                        9.40625  15.875 14.40625],
                      [15.96875  24.625 20.96875;
                       32.875    50.5   42.875;
                       16.90625  25.875 21.90625])
        @fact B --> roughly(Btarget)
        Argb = reinterpret(RGB, reinterpret(Ufixed16, permutedims(A, (3,1,2))))
        B = Images.restrict(Argb)
        Bf = permutedims(reinterpret(Float64, B), (2,3,1))
        @fact Bf --> roughly(Btarget/reinterpret(one(Ufixed16)), 1e-12)
        Argba = reinterpret(RGBA{Ufixed16}, reinterpret(Ufixed16, A))
        B = Images.restrict(Argba)
        @fact reinterpret(Float64, B) --> roughly(Images.restrict(A, (2,3))/reinterpret(one(Ufixed16)), 1e-12)
        A = reshape(1:60, 5, 4, 3)
        B = Images.restrict(A, (1,2,3))
        @fact cat(3, [ 2.6015625  8.71875 6.1171875;
                       4.09375   12.875   8.78125;
                       3.5390625 10.59375 7.0546875],
                     [10.1015625 23.71875 13.6171875;
                      14.09375   32.875   18.78125;
                      11.0390625 25.59375 14.5546875]) --> roughly(B)
        Images.restrict(imgcol, (1,2))
    end

    context("Erode/ dilate") do
        A = zeros(4,4,3)
        A[2,2,1] = 0.8
        A[4,4,2] = 0.6
        Ae = Images.erode(A)
        @fact Ae --> zeros(size(A))
        Ad = Images.dilate(A)
        Ar = [0.8 0.8 0.8 0;
              0.8 0.8 0.8 0;
              0.8 0.8 0.8 0;
              0 0 0 0]
        Ag = [0 0 0 0;
              0 0 0 0;
              0 0 0.6 0.6;
              0 0 0.6 0.6]
        @fact Ad --> cat(3, Ar, Ag, zeros(4,4))
        Ae = Images.erode(Ad)
        Ar = [0.8 0.8 0 0;
              0.8 0.8 0 0;
              0 0 0 0;
              0 0 0 0]
        Ag = [0 0 0 0;
              0 0 0 0;
              0 0 0 0;
              0 0 0 0.6]
        @fact Ae --> cat(3, Ar, Ag, zeros(4,4))
        # issue #311
        @fact Images.dilate(trues(3)) --> trues(3)
    end

    context("Opening / closing") do
        A = zeros(4,4,3)
        A[2,2,1] = 0.8
        A[4,4,2] = 0.6
        Ao = Images.opening(A)
        @fact Ao --> zeros(size(A))
        A = zeros(10,10)
        A[4:7,4:7] = 1
        B = copy(A)
        A[5,5] = 0
        Ac = Images.closing(A)
        @fact Ac --> B
    end

    context("Label components") do
        A = [true  true  false true;
             true  false true  true]
        lbltarget = [1 1 0 2;
                     1 0 2 2]
        lbltarget1 = [1 2 0 4;
                      1 0 3 4]
        @fact Images.label_components(A) --> lbltarget
        @fact Images.label_components(A, [1]) --> lbltarget1
        connectivity = [false true  false;
                        true  false true;
                        false true  false]
        @fact Images.label_components(A, connectivity) --> lbltarget
        connectivity = trues(3,3)
        lbltarget2 = [1 1 0 1;
                      1 0 1 1]
        @fact Images.label_components(A, connectivity) --> lbltarget2
    end

    context("Phantoms") do
        P = [ 0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0;
              0.0  0.0  1.0  0.2  0.2  1.0  0.0  0.0;
              0.0  0.0  0.2  0.3  0.3  0.2  0.0  0.0;
              0.0  0.0  0.2  0.0  0.2  0.2  0.0  0.0;
              0.0  0.0  0.2  0.0  0.0  0.2  0.0  0.0;
              0.0  0.0  0.2  0.2  0.2  0.2  0.0  0.0;
              0.0  0.0  1.0  0.2  0.2  1.0  0.0  0.0;
              0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0 ]
        Q = Images.shepp_logan(8)
        @fact norm((P-Q)[:]) --> less_than(1e-10)
        P = [ 0.0  0.0  0.0   0.0   0.0   0.0   0.0  0.0;
              0.0  0.0  2.0   1.02  1.02  2.0   0.0  0.0;
              0.0  0.0  1.02  1.03  1.03  1.02  0.0  0.0;
              0.0  0.0  1.02  1.0   1.02  1.02  0.0  0.0;
              0.0  0.0  1.02  1.0   1.0   1.02  0.0  0.0;
              0.0  0.0  1.02  1.02  1.02  1.02  0.0  0.0;
              0.0  0.0  2.0   1.02  1.02  2.0   0.0  0.0;
              0.0  0.0  0.0   0.0   0.0   0.0   0.0  0.0 ]
        Q = Images.shepp_logan(8, highContrast=false)
        @fact norm((P-Q)[:]) --> less_than(1e-10)
    end

    context("Image resize") do
        img = convert(Images.Image, zeros(10,10))
        img2 = Images.imresize(img, (5,5))
        @fact length(img2) --> 25
    end

context("Contrast") do
    # Issue #282
    img = convert(Images.Image{Gray{Ufixed8}}, eye(2,2))
    imgs = Images.imstretch(img, 0.3, 0.4)
    @fact data(imgs) --> roughly(1./(1 + (0.3./(eye(2,2) + eps())).^0.4))
end

end


# DiffusionMap.jl
module DiffusionMap

# dependency
using LinearAlgebra, Arpack, KrylovKit, SparseArrays

export
# kernels.jl
AbstractKernel, InverseDistanceKernel,
GaussianKernel, CosineKernel, CustomKernel,

# laplacians.jl
AbstractLaplacian, RowNormalizedLaplacian,
Adjacency, NormalizedAdjacency,
SymmetricLaplacian, RegularLaplacian,

# eigensolvers.jl
AbstractEigenSolver, FullEigen,
ArpackEigen, KrylovEigen,

# diffusionTypes.jl

Diffusionmap,

# diffusionCalculation.jl
get_laplacian, get_adjacency, thresholding!, solve

include("kernels.jl")
include("laplacians.jl")
include("eigensolvers.jl")
include("diffusionTypes.jl")
include("diffusionCalculation.jl")
# include("utils.jl")

end # module


using DataFrames
using SemiparametricAUC

# Data analysis examples
fasd = DataFrames.readtable(joinpath(Pkg.dir("SemiparametricAUC"), "data/fasd.csv"))
# fasd = readtable("ds.csv")

# Define factor/categorical variable
function convert_to_factor(x)
    return(DataFrames.pool(x))
end

fasd[:group] = convert_to_factor(fasd[:group])
fasd[:x1]    = convert_to_factor(fasd[:x1])
fasd[:x2]    = convert_to_factor(fasd[:x2])
# fasd[:x3]  = convert_to_factor(fasd[:x3])

# Perform analysis
one_covariates_results = SemiparametricAUC.semiparametricAUC(model_formula = y ~ x1,
                                              treatment_group = :group, data = fasd)

two_covariates_results = SemiparametricAUC.semiparametricAUC(model_formula = y ~ x1 + x2,
                                            treatment_group = :group, data = fasd)


using Gillespie
import Random: seed!

function F_l(x,parms)
    (Y1,Y2) = x
    (c1,c2,c3) = parms
    [c1*Y1,c2*Y1*Y2,c3*Y2]
end

x0 = [1000,1000]
nu = [[1 0];[-1 1];[0 -1]]
parms = [10.0,0.01,10.0]
tf = 2.0
seed!(1234)

gillespie_result = ssa(x0,F_l,nu,parms,tf)
gillespie_data = ssa_data(gillespie_result)

jensen_result = ssa(x0,F_l,nu,parms,tf,algo=:jensen,max_rate=100000.0,thin=true)
jensen_data = ssa_data(jensen_result)

jensen_alljumps_result = ssa(x0,F_l,nu,parms,tf,algo=:jensen,max_rate=100000.0,thin=false)
jensen_alljumps_data = ssa_data(jensen_alljumps_result)


function neural_ode(model,x,tspan,
                    args...;kwargs...)
  Tracker.istracked(x) && error("u0 is not currently differentiable.")
  p = destructure(model)
  dudt_(du,u::TrackedArray,p,t) = du .= restructure(model,p)(u)
  dudt_(du,u::AbstractArray,p,t) = du .= Flux.data(restructure(model,p)(u))
  prob = ODEProblem(dudt_,x,tspan,p)
  return diffeq_adjoint(p,prob,args...;kwargs...)
end

neural_msde(x,model,mp,tspan,args...;kwargs...) = neural_msde(x,model,mp,tspan,
                                                         diffeq_fd,
                                                         args...;kwargs...)
function neural_msde(model,x,mp,tspan,
                    ad_func::Function,
                    args...;kwargs...)
  p = Flux.data(destructure(model))
  dudt_(du,u::TrackedArray,p,t) = du .= restructure(model,p)(u)
  dudt_(du,u::AbstractArray,p,t) = du .= Flux.data(restructure(model,p)(u))
  g(du,u,p,t) = du .= mp.*u
  prob = SDEProblem(dudt_,g,x,tspan,p)

  if ad_func === diffeq_adjoint
    return ad_func(p,prob,args...;kwargs...)
  elseif ad_func === diffeq_fd
    return ad_func(p,Array,length(p),prob,args...;kwargs...)
  else
    return ad_func(p,Array,prob,args...;kwargs...)
  end
end


using AlmostBandedMatrices, FillArrays, Compat.Test, Compat.LinearAlgebra
using Compat

# Constructions:
@testset "construction" begin
    let n = 10, v = rand(n, 2), w = rand(n, 2), y = rand(n, 3), z = rand(n, 3), U = LowRankMatrix(v, w), L = LowRankMatrix(y, z)
        @test SemiSeparableMatrix(U, L, 1, 2) isa SemiSeparableMatrix{Float64}
        @test Matrix(SemiSeparableMatrix(U, L, 1, 2)) == triu(U, 2)+tril(L, -3)
        @test_throws AssertionError SemiSeparableMatrix(U, L, 10, 2)
        @test_throws AssertionError SemiSeparableMatrix(U, L, 1, 10)
    end

    let v = rand(n, 2), w = rand(n, 2), y = rand(n+1, 3), z = rand(n+1, 3), U = LowRankMatrix(v, w), L = LowRankMatrix(y, z)
        @test_throws AssertionError SemiSeparableMatrix(U, L, 1, 2)
    end
end


module SemiclassicalOrthogonalPolynomials
using ClassicalOrthogonalPolynomials, FillArrays, LazyArrays, ArrayLayouts, QuasiArrays, InfiniteArrays, ContinuumArrays, LinearAlgebra, BandedMatrices,
        SpecialFunctions, HypergeometricFunctions

import Base: getindex, axes, size, \, /, *, +, -, summary, ==, copy, sum, unsafe_getindex, convert, OneTo

import ArrayLayouts: MemoryLayout, ldiv, diagonaldata, subdiagonaldata, supdiagonaldata
import BandedMatrices: bandwidths, AbstractBandedMatrix, BandedLayout, _BandedMatrix
import LazyArrays: resizedata!, paddeddata, CachedVector, CachedMatrix, CachedAbstractVector, LazyMatrix, LazyVector, arguments, ApplyLayout, colsupport, AbstractCachedVector, AccumulateAbstractVector, LazyVector, AbstractCachedMatrix
import ClassicalOrthogonalPolynomials: OrthogonalPolynomial, recurrencecoefficients, jacobimatrix, normalize, _p0, UnitInterval, orthogonalityweight, NormalizedBasisLayout,
                                        Bidiagonal, Tridiagonal, SymTridiagonal, symtridiagonalize, normalizationconstant, LanczosPolynomial,
                                        OrthogonalPolynomialRatio, Weighted, Expansion, UnionDomain, oneto, Hilbert, WeightedOrthogonalPolynomial, HalfWeighted,
                                        Associated
import InfiniteArrays: OneToInf, InfUnitRange
import ContinuumArrays: basis, Weight, @simplify, AbstractBasisLayout, BasisLayout, MappedBasisLayout, grid
import FillArrays: SquareEye
import HypergeometricFunctions: _₂F₁general2

export LanczosPolynomial, Legendre, Normalized, normalize, SemiclassicalJacobi, SemiclassicalJacobiWeight, WeightedSemiclassicalJacobi, OrthogonalPolynomialRatio, TwoBandJacobi, TwoBandWeight

""""
    SemiclassicalJacobiWeight(t, a, b, c)

is a quasi-vector corresponding to the weight `x^a * (1-x)^b * (t-x)^c` on `x in 0..1`
"""
struct SemiclassicalJacobiWeight{T} <: Weight{T}
    t::T
    a::T
    b::T
    c::T
end

function SemiclassicalJacobiWeight(t, a, b, c)
    T = promote_type(eltype(t), eltype(a), eltype(b), eltype(c))
    SemiclassicalJacobiWeight{T}(t, a, b, c)
end

copy(w::SemiclassicalJacobiWeight) = w

axes(P::SemiclassicalJacobiWeight{T}) where T = (Inclusion(UnitInterval{T}()),)
function getindex(P::SemiclassicalJacobiWeight, x::Real)
    t,a,b,c = P.t,P.a,P.b,P.c
    @boundscheck checkbounds(P, x)
    x^a * (1-x)^b * (t-x)^c
end

function sum(P::SemiclassicalJacobiWeight{T}) where T
    (t,a,b,c) = BigFloat.((P.t,P.a,P.b,P.c))
    return convert(T, t^c * beta(1+a,1+b) * _₂F₁general2(1+a,-c,2+a+b,1/t))
end

function summary(io::IO, P::SemiclassicalJacobiWeight)
    t,a,b,c = P.t,P.a,P.b,P.c
    print(io, "x^$a * (1-x)^$b * ($t-x)^$c")
end

function ==(A::SemiclassicalJacobiWeight, B::SemiclassicalJacobiWeight)
    A.a == B.a && A.b == B.b &&
        ((iszero(A.c) && iszero(B.c)) ||
         (A.t == B.t && A.c == B.c))
end

function Expansion(w::SemiclassicalJacobiWeight{T}) where T
    t,a,b,c = w.t,w.a,w.b,w.c
    P = jacobi(b, a, UnitInterval{T}())
    x = axes(P,1)
    # TODO: simplify
    LanczosPolynomial(@.(x^a * (1-x)^b * (t-x)^c), P).w
end

==(A::SemiclassicalJacobiWeight, B::Expansion) = Expansion(A) == B
==(A::Expansion, B::SemiclassicalJacobiWeight) = A == Expansion(B)

"""
   RaisedOP(P, y)

gives the OPs w.r.t. (y - x) .* w based on lowering to Q.
"""

struct RaisedOP{T, QQ, LL<:OrthogonalPolynomialRatio} <: OrthogonalPolynomial{T}
    Q::QQ
    ℓ::LL
end

RaisedOP(Q, ℓ::OrthogonalPolynomialRatio) = RaisedOP{eltype(Q),typeof(Q),typeof(ℓ)}(Q, ℓ)
RaisedOP(Q, y::Number) = RaisedOP(Q, OrthogonalPolynomialRatio(Q,y))

function jacobimatrix(P::RaisedOP{T}) where T
    ℓ = P.ℓ
    X = jacobimatrix(P.Q)
    a,b = diagonaldata(X), supdiagonaldata(X)
    # non-normalized lower diag of Jacobi
    v = Vcat(zero(T),b .* ℓ)
    c = BroadcastVector((ℓ,a,b,sa,v) -> ℓ*a + b - b*ℓ^2 - sa*ℓ + ℓ*v, ℓ, a, b, a[2:∞], v)
    Tridiagonal(c, BroadcastVector((ℓ,a,b,v) -> a - b * ℓ + v, ℓ,a,b,v), b)
end

"""
  the bands of the Jacobi matrix
"""
mutable struct SemiclassicalJacobiBand{dv,T} <: AbstractCachedVector{T}
    data::Vector{T}
    a::AbstractVector{T}
    b::AbstractVector{T}
    ℓ::AbstractVector{T}
    datasize::Tuple{Int}
end

function LazyArrays.cache_filldata!(r::SemiclassicalJacobiBand{:dv}, inds::AbstractUnitRange)
    rℓ = r.ℓ[inds[1]-1:inds[end]]; ℓ,sℓ = rℓ[2:end], rℓ[1:end-1]
    ra = r.a[inds[1]:(inds[end]+1)]; a = ra[1:end-1]; sa = ra[2:end]
    rb = r.b[inds[1]-1:inds[end]]; b = rb[2:end]; sb = rb[1:end-1]
    r.data[inds] .= @.(a - b * ℓ + sb*sℓ)
end

function LazyArrays.cache_filldata!(r::SemiclassicalJacobiBand{:ev}, inds::AbstractUnitRange)
    rℓ = r.ℓ[inds[1]-1:inds[end]]; ℓ,sℓ = rℓ[2:end], rℓ[1:end-1]
    ra = r.a[inds[1]:(inds[end]+1)]; a = ra[1:end-1]; sa = ra[2:end]
    rb = r.b[inds[1]-1:inds[end]]; b = rb[2:end]; sb = rb[1:end-1]
    r.data[inds] .= @.(sqrt((ℓ*a + b - b*ℓ^2 - sa*ℓ + ℓ*sb*sℓ)*b))
end

size(::SemiclassicalJacobiBand) = (ℵ₀,)

SemiclassicalJacobiBand{:dv,T}(a,b,ℓ) where T = SemiclassicalJacobiBand{:dv,T}(T[a[1] - b[1]ℓ[1]], a, b, ℓ, (1,))
SemiclassicalJacobiBand{:ev,T}(a,b,ℓ) where T = SemiclassicalJacobiBand{:ev,T}(T[sqrt((ℓ[1]*a[1] + b[1] - b[1]*ℓ[1]^2 - a[2]*ℓ[1])b[1])], a, b, ℓ, (1,))
SemiclassicalJacobiBand{dv}(a,b,ℓ) where dv = SemiclassicalJacobiBand{dv,promote_type(eltype(a),eltype(b),eltype(ℓ))}(a,b,ℓ)

copy(r::SemiclassicalJacobiBand) = r # immutable

"""
   SemiclassicalJacobi(t, a, b, c)

is a quasi-matrix for the  orthogonal polynomials w.r.t. x^a * (1-x)^b * (t-x)^c on [0,1]
"""
struct SemiclassicalJacobi{T} <: OrthogonalPolynomial{T}
    t::T
    a::T
    b::T
    c::T
    X::AbstractMatrix{T}
    SemiclassicalJacobi{T}(t::T,a::T,b::T,c::T,X::AbstractMatrix{T}) where T = new{T}(t,a,b,c,X)
end

const WeightedSemiclassicalJacobi{T} = WeightedBasis{T,<:SemiclassicalJacobiWeight,<:SemiclassicalJacobi}

function SemiclassicalJacobi(t, a, b, c, X::AbstractMatrix)
    T = float(promote_type(typeof(t), typeof(a), typeof(b), typeof(c), eltype(X)))
    SemiclassicalJacobi{T}(T(t), T(a), T(b), T(c), convert(AbstractMatrix{T},X))
end

SemiclassicalJacobi(t, a, b, c, P::SemiclassicalJacobi) = SemiclassicalJacobi(t, a, b, c, semiclassical_jacobimatrix(P, a, b, c))
SemiclassicalJacobi(t, a, b, c) = SemiclassicalJacobi(t, a, b, c, semiclassical_jacobimatrix(t, a, b, c))
SemiclassicalJacobi{T}(t, a, b, c) where T = SemiclassicalJacobi(convert(T,t), convert(T,a), convert(T,b), convert(T,b))

WeightedSemiclassicalJacobi(t, a, b, c, P...) = SemiclassicalJacobiWeight(t, a, b, c) .* SemiclassicalJacobi(t, a, b, c, P...)

function semiclassical_jacobimatrix(t, a, b, c)
    T = float(promote_type(typeof(t), typeof(a), typeof(b), typeof(c)))
    if a == 0 && b == 0 && c == -1
        # for this special case we can generate the Jacobi operator explicitly
        N = (1:∞)
        α = T.(neg1c_αcfs(t)) # cached coefficients
        A = Vcat((α[1]+1)/2 , -N./(N.*4 .- 2).*α .+ (N.+1)./(N.*4 .+ 2).*α[2:end].+1/2)
        C = -(N)./(N.*4 .- 2)
        B = Vcat((α[1]^2*3-α[1]*α[2]*2-1)/6 , -(N)./(N.*4 .+ 2).*α[2:end]./α)
        # if J is Tridiagonal(c,a,b) then for norm. OPs it becomes SymTridiagonal(a, sqrt.( b.* c))
        return SymTridiagonal(A, sqrt.(B.*C))
    end
    P = jacobi(b, a, UnitInterval{T}())
    iszero(c) && return symtridiagonalize(jacobimatrix(P))
    x = axes(P,1)
    jacobimatrix(LanczosPolynomial(@.(x^a * (1-x)^b * (t-x)^c), P))
end

function symraised_jacobimatrix(Q, y)
    ℓ = OrthogonalPolynomialRatio(Q,y)
    X = jacobimatrix(Q)
    a,b = diagonaldata(X), supdiagonaldata(X)
    SymTridiagonal(SemiclassicalJacobiBand{:dv}(a,b,ℓ), SemiclassicalJacobiBand{:ev}(a,b,ℓ))
end

function semiclassical_jacobimatrix(Q::SemiclassicalJacobi, a, b, c)
    if  iszero(a) && iszero(b) && c == -1 # (a,b,c) = (0,0,-1) special case
        semiclassical_jacobimatrix(Q.t, a, b, c)
    elseif a == Q.a+1 && b == Q.b && c == Q.c  # raising by 1
        symraised_jacobimatrix(Q, 0)
    elseif a == Q.a && b == Q.b+1 && c == Q.c
        symraised_jacobimatrix(Q, 1)
    elseif a == Q.a && b == Q.b && c == Q.c+1
        symraised_jacobimatrix(Q, Q.t)
    elseif a == Q.a && b == Q.b && c == Q.c-1 # lowering by 1
        symlowered_jacobimatrix(Q, :c)
    elseif a == Q.a-1 && b == Q.b && c == Q.c
        symlowered_jacobimatrix(Q, :a)
    elseif a == Q.a && b == Q.b-1 && c == Q.c
        symlowered_jacobimatrix(Q, :b)
    elseif a > Q.a  # iterative raising
        semiclassical_jacobimatrix(SemiclassicalJacobi(Q.t, Q.a+1, Q.b, Q.c, Q), a, b, c)
    elseif b > Q.b
        semiclassical_jacobimatrix(SemiclassicalJacobi(Q.t, Q.a, Q.b+1, Q.c, Q), a, b, c)
    elseif c > Q.c
        semiclassical_jacobimatrix(SemiclassicalJacobi(Q.t, Q.a, Q.b, Q.c+1, Q), a, b, c)
    elseif b < Q.b  # iterative lowering
        semiclassical_jacobimatrix(SemiclassicalJacobi(Q.t, Q.a, Q.b-1, Q.c, Q), a, b, c)
    elseif c < Q.c
        semiclassical_jacobimatrix(SemiclassicalJacobi(Q.t, Q.a, Q.b, Q.c-1, Q), a, b, c)
    elseif a < Q.a
        semiclassical_jacobimatrix(SemiclassicalJacobi(Q.t, Q.a-1, Q.b, Q.c, Q), a, b, c)
    elseif a == Q.a && b == Q.b && c == Q.c # same basis
        jacobimatrix(Q)
    else
        error("Not Implemented")
    end
end

LanczosPolynomial(P::SemiclassicalJacobi{T}) where T =
    LanczosPolynomial(orthogonalityweight(P), Normalized(jacobi(P.b, P.a, UnitInterval{T}())), P.X.dv.data)

"""
    toclassical(P::SemiclassicalJacobi)

gives either a mapped `Jacobi` or `LanczosPolynomial` version of `P`.
"""
toclassical(P::SemiclassicalJacobi{T}) where T = iszero(P.c) ? Normalized(jacobi(P.b, P.a, UnitInterval{T}())) : LanczosPolynomial(P)

copy(P::SemiclassicalJacobi) = P
axes(P::SemiclassicalJacobi{T}) where T = (Inclusion(UnitInterval{T}()),OneToInf())

==(A::SemiclassicalJacobi, B::SemiclassicalJacobi) = A.t == B.t && A.a == B.a && A.b == B.b && A.c == B.c
==(::AbstractQuasiMatrix, ::SemiclassicalJacobi) = false
==(::SemiclassicalJacobi, ::AbstractQuasiMatrix) = false

orthogonalityweight(P::SemiclassicalJacobi) = SemiclassicalJacobiWeight(P.t, P.a, P.b, P.c)

function summary(io::IO, P::SemiclassicalJacobi)
    t,a,b,c = P.t,P.a,P.b,P.c
    print(io, "SemiclassicalJacobi with weight x^$a * (1-x)^$b * ($t-x)^$c")
end

jacobimatrix(P::SemiclassicalJacobi) = P.X

"""
    op_lowering(Q, y)
Gives the Lowering operator from OPs w.r.t. (x-y)*w(x) to Q
as constructed from Chistoffel–Darboux
"""
function op_lowering(Q, y)
    # we first use Christoff-Darboux with d = 1
    # But we want the first OP to be 1 so we rescale
    P = RaisedOP(Q, y)
    A,_,_ = recurrencecoefficients(Q)
    d = -inv(A[1]*_p0(Q)*P.ℓ[1])
    κ = d * normalizationconstant(1, P)
    T = eltype(κ)
    # hide array type for compilation
    Bidiagonal(κ, -(κ .* P.ℓ), :L)
end

function semijacobi_ldiv(Q, P::SemiclassicalJacobi)
    if P.a ≤ 0 && P.b ≤ 0 && P.c ≤ 0
        P̃ = toclassical(P)
        (Q \ P̃)/_p0(P̃)
    else
        error("Implement")
    end
end

function semijacobi_ldiv(P::SemiclassicalJacobi, Q)
    R = SemiclassicalJacobi(P.t, mod(P.a,-1), mod(P.b,-1), mod(P.c,-1))
    R̃ = toclassical(R)
    (P \ R) * _p0(R̃) * (R̃ \ Q)
end

struct SemiclassicalJacobiLayout <: AbstractBasisLayout end
MemoryLayout(::Type{<:SemiclassicalJacobi}) = SemiclassicalJacobiLayout()

copy(L::Ldiv{<:NormalizedBasisLayout,SemiclassicalJacobiLayout}) = copy(Ldiv{ApplyLayout{typeof(*)},SemiclassicalJacobiLayout}(L.A, L.B))
copy(L::Ldiv{SemiclassicalJacobiLayout,<:NormalizedBasisLayout}) = copy(Ldiv{SemiclassicalJacobiLayout,ApplyLayout{typeof(*)}}(L.A, L.B))

copy(L::Ldiv{ApplyLayout{typeof(*)},SemiclassicalJacobiLayout}) = copy(Ldiv{ApplyLayout{typeof(*)},BasisLayout}(L.A, L.B))
copy(L::Ldiv{SemiclassicalJacobiLayout,ApplyLayout{typeof(*)}}) = copy(Ldiv{BasisLayout,ApplyLayout{typeof(*)}}(L.A, L.B))

copy(L::Ldiv{MappedBasisLayout,SemiclassicalJacobiLayout}) = semijacobi_ldiv(L.A, L.B)
copy(L::Ldiv{SemiclassicalJacobiLayout,MappedBasisLayout}) = semijacobi_ldiv(L.A, L.B)

copy(L::Ldiv{SemiclassicalJacobiLayout}) = semijacobi_ldiv(L.A, L.B)
copy(L::Ldiv{<:Any,SemiclassicalJacobiLayout}) = semijacobi_ldiv(L.A, L.B)
function copy(L::Ldiv{SemiclassicalJacobiLayout,SemiclassicalJacobiLayout})
    Q,P = L.A,L.B
    @assert Q.t == P.t
    Q == P && return SquareEye{eltype(L)}((axes(P,2),))
    M_Q = massmatrix(Q)
    M_P = massmatrix(P)
    L = P \ (SemiclassicalJacobiWeight(Q.t, Q.a-P.a, Q.b-P.b, Q.c-P.c) .* Q)
    inv(M_Q) * L' * M_P
end

\(A::LanczosPolynomial, B::SemiclassicalJacobi) = semijacobi_ldiv(A, B)
\(A::SemiclassicalJacobi, B::LanczosPolynomial) = semijacobi_ldiv(A, B)
function \(w_A::WeightedSemiclassicalJacobi, w_B::WeightedSemiclassicalJacobi)
    wA,A = w_A.args
    wB,B = w_B.args
    @assert wA.t == wB.t == A.t == B.t

    if wA.a == wB.a && wA.b == wB.b && wA.c == wB.c
        A \ B
    elseif wA.a+1 == wB.a && wA.b == wB.b && wA.c == wB.c
        @assert A.a+1 == B.a && A.b == B.b && A.c == B.c
        op_lowering(A,0)
    elseif wA.a == wB.a && wA.b+1 == wB.b && wA.c == wB.c
        @assert A.a == B.a && A.b+1 == B.b && A.c == B.c
        -op_lowering(A,1)
    elseif wA.a == wB.a && wA.b == wB.b && wA.c+1 == wB.c
        @assert A.a == B.a && A.b == B.b && A.c+1 == B.c
        -op_lowering(A,A.t)
    elseif wA.a+1 ≤ wB.a
        C = SemiclassicalJacobi(A.t, A.a+1, A.b, A.c, A)
        w_C = SemiclassicalJacobiWeight(wA.t, wA.a+1, wA.b, wA.c) .* C
        L_2 = w_C \ w_B
        L_1 = w_A \ w_C
        L_1 * L_2
    elseif wA.b+1 ≤ wB.b
        C = SemiclassicalJacobi(A.t, A.a, A.b+1, A.c, A)
        w_C = SemiclassicalJacobiWeight(wA.t, wA.a, wA.b+1, wA.c) .* C
        L_2 = w_C \ w_B
        L_1 = w_A \ w_C
        L_1 * L_2
    else
        error("Implement")
    end
end

\(A::SemiclassicalJacobi, w_B::WeightedSemiclassicalJacobi) = (SemiclassicalJacobiWeight(A.t,0,0,0) .* A) \ w_B

massmatrix(P::SemiclassicalJacobi) = Diagonal(Fill(sum(orthogonalityweight(P)),∞))

@simplify function *(Ac::QuasiAdjoint{<:Any,<:SemiclassicalJacobi}, wB::WeightedBasis{<:Any,<:SemiclassicalJacobiWeight,<:SemiclassicalJacobi})
    A = parent(Ac)
    w,B = arguments(wB)
    P = SemiclassicalJacobi(w.t, w.a, w.b, w.c)
    (P\A)' * massmatrix(P) * (P \ B)
end

function ldiv(Q::SemiclassicalJacobi, f::AbstractQuasiVector)
    if iszero(Q.a) && iszero(Q.b) && Q.c == -1
        # todo: due to a stdlib error this won't work with bigfloat as is
        T = typeof(Q.t)
        R = Legendre{T}()[affine(Inclusion(zero(T)..one(T)), axes(Legendre{T}(),1)), :]
        B = neg1c_tolegendre(Q.t)
        return (B \ (R \ f))
    end
    R = SemiclassicalJacobi(Q.t, mod(Q.a,-1), mod(Q.b,-1), mod(Q.c,-1))
    R̃ = toclassical(R)
    return (Q \ R̃) * (R̃ \ f)
end
function ldiv(Qn::SubQuasiArray{<:Any,2,<:SemiclassicalJacobi,<:Tuple{<:Inclusion,<:Any}}, C::AbstractQuasiArray)
    _,jr = parentindices(Qn)
    Q = parent(Qn)
    if iszero(Q.a) && iszero(Q.b) && Q.c == -1
        # todo: due to a stdlib error this won't work with bigfloat as is
        T = typeof(Q.t)
        R = Legendre{T}()[affine(Inclusion(zero(T)..one(T)), axes(Legendre{T}(),1)), :]
        B = neg1c_tolegendre(Q.t)
        return (B[jr,jr] \ (R[:,jr] \ C))
    end
    R = SemiclassicalJacobi(Q.t, mod(Q.a,-1), mod(Q.b,-1), mod(Q.c,-1))
    R̃ = toclassical(R)
    (Q \ R̃)[jr,jr] * (R̃[:,jr] \ C)
end

# sqrt(1-(1-x)^2) == sqrt(2x-x^2) == sqrt(x)*sqrt(2-x)
# sqrt(1-(1-x)^2) == sqrt(2x-x^2) == sqrt(x)*sqrt(2-x)

convert(::Type{WeightedOrthogonalPolynomial}, Q::HalfWeighted{:a,T,<:SemiclassicalJacobi}) where T = SemiclassicalJacobiWeight(Q.P.t, Q.P.a,zero(T),zero(T)) .* Q.P
convert(::Type{WeightedOrthogonalPolynomial}, Q::HalfWeighted{:b,T,<:SemiclassicalJacobi}) where T = SemiclassicalJacobiWeight(Q.P.t, zero(T),Q.P.b,zero(T)) .* Q.P
convert(::Type{WeightedOrthogonalPolynomial}, Q::HalfWeighted{:c,T,<:SemiclassicalJacobi}) where T = SemiclassicalJacobiWeight(Q.P.t, zero(T),zero(T),Q.P.c) .* Q.P

include("derivatives.jl")

###
# Hierarchy
###

function Base.broadcasted(::Type{SemiclassicalJacobi}, t::Number, ar::AbstractUnitRange, b::Number, c::Number)
    Ps = [SemiclassicalJacobi(t, first(ar), b, c)]
    for a in ar[2:end]
        push!(Ps, SemiclassicalJacobi(t, a, b, c, Ps[end]))
    end
    Ps
end

function Base.broadcasted(::Type{SemiclassicalJacobi}, t::Number, a::Number, br::AbstractUnitRange, c::Number)
    Ps = [SemiclassicalJacobi(t, a, first(br), c)]
    for b in br[2:end]
        push!(Ps, SemiclassicalJacobi(t, a, b, c, Ps[end]))
    end
    Ps
end

function Base.broadcasted(::Type{SemiclassicalJacobi}, t::Number, a::Number, b::Number, cr::AbstractUnitRange)
    Ps = [SemiclassicalJacobi(t, a, b, first(cr))]
    for c in cr[2:end]
        push!(Ps, SemiclassicalJacobi(t, a, b, c, Ps[end]))
    end
    Ps
end

include("twoband.jl")
include("lowering.jl")

end


"""
    test_bqg_rossbywave(; kwargs...)

Evolves a Rossby wave and compares with the analytic solution.
"""
function test_bqg_rossbywave(stepper, dt, nsteps)
    nx = 64
  beta = 2.0
    Lx = 2π
    mu = 0.0
    nu = 0.0

  gr  = TwoDGrid(nx, Lx)
  x, y = gridpoints(gr)

  # the following if statement is called so that all the cases of
  # Problem() fuction are tested
  if stepper=="ForwardEuler"
    eta = zeros(nx, nx)
  else
    eta(x, y) = 0*x
  end

  prob = BarotropicQG.InitialValueProblem(nx=nx, Lx=Lx, eta=eta, beta=beta, mu=mu, nu=nu, 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 = 3.0*2π/g.Lx
  lwave = 2.0*2π/g.Ly
      ω = -p.beta*kwave/(kwave^2.0 + lwave^2.0)
     ζ0 = @. ampl*cos(kwave*x)*cos(lwave*y)
    ζ0h = rfft(ζ0)

  BarotropicQG.set_zeta!(prob, ζ0)

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

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

  isapprox(ζ_theory, v.zeta, rtol=g.nx*g.ny*nsteps*1e-12)
end

"""
    test_stochasticforcingbudgets(; kwargs...)

Tests if the energy budgets are closed for BarotropicQG with stochastic forcing.
"""
function test_bqg_stochasticforcingbudgets(; n=256, dt=0.01, L=2π, nu=1e-7, nnu=2, mu=1e-1, message=false)
  n, L  = 256, 2π
  nu, nnu = 1e-7, 2
  mu = 1e-1
  dt, tf = 0.005, 0.1/mu
  nt = round(Int, tf/dt)
  ns = 1

  # Forcing
  kf, dkf = 12.0, 2.0
  ε = 0.1

  gr  = TwoDGrid(n, L)
  x, y = gridpoints(gr)

  Kr = [ gr.kr[i] for i=1:gr.nkr, j=1:gr.nl]

  force2k = zero(gr.Krsq)
  @. force2k = exp.(-(sqrt(gr.Krsq)-kf)^2/(2*dkf^2))
  @. force2k[gr.Krsq .< 2.0^2 ] = 0
  @. force2k[gr.Krsq .> 20.0^2 ] = 0
  @. force2k[Kr .< 2π/L] = 0
  ε0 = parsevalsum(force2k.*gr.invKrsq/2.0, gr)/(gr.Lx*gr.Ly)
  force2k .= ε/ε0 * force2k

  Random.seed!(1234)

  function calcFq!(Fqh, sol, t, cl, v, p, g)
    eta = exp.(2π*im*rand(Float64, size(sol)))/sqrt(cl.dt)
    eta[1, 1] = 0
    @. Fqh = eta * sqrt(force2k)
    nothing
  end

  prob = BarotropicQG.ForcedProblem(nx=n, Lx=L, nu=nu, nnu=nnu, mu=mu, dt=dt,
   stepper="RK4", calcFq=calcFq!, stochastic=true)

  sol, cl, v, p, g = prob.sol, prob.clock, prob.vars, prob.params, prob.grid

  BarotropicQG.set_zeta!(prob, 0*x)
  E = Diagnostic(BarotropicQG.energy,      prob, nsteps=nt)
  D = Diagnostic(BarotropicQG.dissipation, prob, nsteps=nt)
  R = Diagnostic(BarotropicQG.drag,        prob, nsteps=nt)
  W = Diagnostic(BarotropicQG.work,        prob, nsteps=nt)
  diags = [E, D, W, R]

  # Step forward

  stepforward!(prob, diags, round(Int, nt))

  BarotropicQG.updatevars!(prob)

  cfl = cl.dt*maximum([maximum(v.v)/g.dx, maximum(v.u)/g.dy])

  E, D, W, R = diags

  t = round(mu*cl.t, digits=2)

  i₀ = 1
  dEdt = (E[(i₀+1):E.i] - E[i₀:E.i-1])/cl.dt
  ii = (i₀):E.i-1
  ii2 = (i₀+1):E.i

  # dEdt = W - D - R?
  # If the Ito interpretation was used for the work
  # then we need to add the drift term
  # total = W[ii2]+ε - D[ii] - R[ii]      # Ito
  total = W[ii2] - D[ii] - R[ii]        # Stratonovich

  residual = dEdt - total

  if message
    println("step: %04d, t: %.1f, cfl: %.3f, time: %.2f s\n", cl.step, cl.t, cfl, t)
  end
  # println(mean(abs.(residual)))
  isapprox(mean(abs.(residual)), 0, atol=1e-4)
end

"""
    test_stochasticforcingbudgets(; kwargs...)

Tests if the energy budgets are closed for BarotropicQG with stochastic forcing.
"""
function test_bqg_deterministicforcingbudgets(; n=256, dt=0.01, L=2π, nu=1e-7, nnu=2, mu=1e-1, message=false)
  n, L  = 256, 2π
  nu, nnu = 1e-7, 2
  mu = 1e-1
  dt, tf = 0.005, 0.1/mu
  nt = round(Int, tf/dt)
  ns = 1

  # Forcing = 0.01cos(4x)cos(5y)cos(2t)
  gr  = TwoDGrid(n, L)
  x, y = gridpoints(gr)
  k0, l0 = gr.kr[2], gr.l[2]

  f = @. 0.01*cos(4*k0*x)*cos(5*l0*y)
  fh = rfft(f)
  function calcFq!(Fqh, sol, t, cl, v, p, g)
    @. Fqh = fh*cos(2*t)
    nothing
  end

  prob = BarotropicQG.ForcedProblem(nx=n, Lx=L, nu=nu, nnu=nnu, mu=mu, dt=dt,
   stepper="RK4", calcFq=calcFq!, stochastic=false)

  sol, cl, v, p, g = prob.sol, prob.clock, prob.vars, prob.params, prob.grid

  BarotropicQG.set_zeta!(prob, 0*x)
  E = Diagnostic(BarotropicQG.energy,      prob, nsteps=nt)
  D = Diagnostic(BarotropicQG.dissipation, prob, nsteps=nt)
  R = Diagnostic(BarotropicQG.drag,        prob, nsteps=nt)
  W = Diagnostic(BarotropicQG.work,        prob, nsteps=nt)
  diags = [E, D, W, R]

  # Step forward

  stepforward!(prob, diags, round(Int, nt))

  BarotropicQG.updatevars!(prob)

  cfl = cl.dt*maximum([maximum(v.v)/g.dx, maximum(v.u)/g.dy])

  E, D, W, R = diags

  t = round(mu*cl.t, digits=2)

  i₀ = 1
  dEdt = (E[(i₀+1):E.i] - E[i₀:E.i-1])/cl.dt
  ii = (i₀):E.i-1
  ii2 = (i₀+1):E.i

  # dEdt = W - D - R?
  total = W[ii2] - D[ii] - R[ii]

  residual = dEdt - total

  if message
    println("step: %04d, t: %.1f, cfl: %.3f, time: %.2f s\n", prob.step, cl.t, cfl, tc)
  end
  # println(mean(abs.(residual)))
  isapprox(mean(abs.(residual)), 0, atol=1e-8)
end

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

Tests the advection term in the twodturb module by timestepping a
test problem with timestep dt and timestepper identified by the string stepper.
The test problem is derived by picking a solution ζf (with associated
streamfunction ψf) for which the advection term J(ψf, ζf) is non-zero. Next, a
forcing Ff is derived according to Ff = ∂ζf/∂t + J(ψf, ζf) - nuΔζf. One solution
to the vorticity equation forced by this Ff is then ζf. (This solution may not
be realized, at least at long times, if it is unstable.)
"""
function test_bqg_advection(dt, stepper; n=128, L=2π, nu=1e-2, nnu=1, mu=0.0, message=false)
  n, L  = 128, 2π
  nu, nnu = 1e-2, 1
  mu = 0.0
  tf = 1.0
  nt = round(Int, tf/dt)

  gr  = TwoDGrid(n, L)
  x, y = gridpoints(gr)

  psif = @. sin(2x)*cos(2y) + 2sin(x)*cos(3y)
  qf = @. -8sin(2x)*cos(2y) - 20sin(x)*cos(3y)

  Ff = @. -(
    nu*( 64sin(2x)*cos(2y) + 200sin(x)*cos(3y) )
    + 8*( cos(x)*cos(3y)*sin(2x)*sin(2y) - 3cos(2x)*cos(2y)*sin(x)*sin(3y) )
  )

  Ffh = rfft(Ff)

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

  prob = BarotropicQG.ForcedProblem(nx=n, Lx=L, nu=nu, nnu=nnu, mu=mu, dt=dt, stepper=stepper, calcFq=calcFq!)
  sol, cl, v, p, g = prob.sol, prob.clock, prob.vars, prob.params, prob.grid
  BarotropicQG.set_zeta!(prob, qf)

  # Step forward
  stepforward!(prob, round(Int, nt))
  BarotropicQG.updatevars!(prob)
  isapprox(v.q, qf, rtol=1e-13)
end

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

Tests the form stress term that forces the domain-averaged zonal flow U(t).
"""
function test_bqg_formstress(dt, stepper; n=128, L=2π, nu=0.0, nnu=1, mu=0.0, message=false)
  n, L  = 128, 2π
  nu, nnu = 1e-2, 1
  mu = 0.0
  tf = 1
  nt = 1

  gr  = TwoDGrid(n, L)
  x, y = gridpoints(gr)

  zetai = @. -20*sin(10*x)*cos(10*y)
  topoPV(x, y) = @. cos(10x)*cos(10y)
  F(t) = 0 #no forcing

  answer = 0.25 # this is what <v*eta> should be

  prob = BarotropicQG.ForcedProblem(nx=n, Lx=L, nu=nu, nnu=nnu, mu=mu, dt=dt, stepper=stepper, eta=topoPV, calcFU = F)
  BarotropicQG.set_zeta!(prob, zetai)
  BarotropicQG.updatevars!(prob)

  # Step forward
  stepforward!(prob, nt)
  isapprox(prob.timestepper.N[1, 1], answer, rtol=1e-13)
end

function test_bqg_energyenstrophy()
  nx, Lx  = 64, 2π
  ny, Ly  = 64, 3π
  g  = TwoDGrid(nx, Lx, ny, Ly)
  k0, l0 = g.k[2], g.l[2] # fundamental wavenumbers
  x, y = gridpoints(g)

  energy_calc = 29/9
  enstrophy_calc = 2701/162

    eta = @. cos(10k0*x)*cos(10l0*y)
   psi0 = @. sin(2k0*x)*cos(2l0*y) + 2sin(k0*x)*cos(3l0*y)
  zeta0 = @. -((2k0)^2+(2l0)^2)*sin(2k0*x)*cos(2l0*y) - (k0^2+(3l0)^2)*2sin(k0*x)*cos(3l0*y)

  prob = BarotropicQG.InitialValueProblem(nx=nx, Lx=Lx, ny=ny, Ly=Ly, eta = eta, stepper="ForwardEuler")
  sol, cl, v, p, g = prob.sol, prob.clock, prob.vars, prob.params, prob.grid
  BarotropicQG.set_zeta!(prob, zeta0)
  BarotropicQG.updatevars!(prob)

  energyzeta0 = BarotropicQG.energy(prob)
  enstrophyzeta0 = BarotropicQG.enstrophy(prob)

  isapprox(energyzeta0, energy_calc, rtol=1e-13) && isapprox(enstrophyzeta0, enstrophy_calc, rtol=1e-13) &&
  BarotropicQG.addforcing!(prob.timestepper.N, sol, cl.t, cl, v, p, g)==nothing
end

function test_bqg_meanenergyenstrophy()
  nx, Lx  = 64, 2π
  ny, Ly  = 96, 3π
  g  = TwoDGrid(nx, Lx, ny, Ly)
  k0, l0 = g.k[2], g.l[2] # fundamental wavenumbers
  x, y = gridpoints(g)

  calcFU(t) = 0.0
  eta(x, y) = @. cos(10x)*cos(10y)
  psi0 = @. sin(2k0*x)*cos(2l0*y) + 2sin(k0*x)*cos(3l0*y)
 zeta0 = @. -((2k0)^2+(2l0)^2)*sin(2k0*x)*cos(2l0*y) - (k0^2+(3l0)^2)*2sin(k0*x)*cos(3l0*y)
  beta = 10.0
  U = 1.2

  energy_calc = 29/9
  enstrophy_calc = 2701/162

  prob = BarotropicQG.ForcedProblem(nx=nx, Lx=Lx, ny=ny, Ly=Ly, beta=beta, eta=eta, calcFU = calcFU,
                                    stepper="ForwardEuler")

  BarotropicQG.set_zeta!(prob, zeta0)
  BarotropicQG.set_U!(prob, U)
  BarotropicQG.updatevars!(prob)

  energyU = BarotropicQG.meanenergy(prob)
  enstrophyU = BarotropicQG.meanenstrophy(prob)

  energyzeta0 = BarotropicQG.energy(prob)
  enstrophyzeta0 = BarotropicQG.enstrophy(prob)

  (isapprox(energyU, 0.5*U^2, rtol=1e-13) &&
    isapprox(enstrophyU, beta*U, rtol=1e-13) &&
    isapprox(energyzeta0, energy_calc, rtol=1e-13) &&
    isapprox(enstrophyzeta0, enstrophy_calc, rtol=1e-13)
  )
end

function test_bqg_problemtype(T=Float32)
  prob = BarotropicQG.Problem(T=T)

  (typeof(prob.sol)==Array{Complex{T},2} && typeof(prob.grid.Lx)==T && typeof(prob.grid.x)==Array{T,2} && typeof(prob.vars.u)==Array{T,2})
end