struct Quaternion{T<:Real} <: Number
    s::T
    v1::T
    v2::T
    v3::T
    norm::Bool
end

(::Type{Quaternion{T}})(x::Real) where {T} = Quaternion{T}(x, 0, 0, 0, false)
(::Type{Quaternion{T}})(q::Quaternion{T}) where {T<:Real} = q
(::Type{Quaternion{T}})(q::Quaternion) where {T<:Real} = Quaternion{T}(q.s, q.v1, q.v2, q.v3, q.norm)
Quaternion(s::Real, v1::Real, v2::Real, v3::Real, n::Bool = false) =
    Quaternion(promote(s, v1, v2, v3)..., n)
Quaternion(x::Real) = Quaternion(x, zero(x), zero(x), zero(x), abs(x) == one(x))
Quaternion(z::Complex) = Quaternion(z.re, z.im, zero(z.re), zero(z.re), abs(z) == one(z.re))
Quaternion(s::Real, a::Vector) = Quaternion(s, a[1], a[2], a[3])
Quaternion(a::Vector) = Quaternion(0, a[1], a[2], a[3])

convert(::Type{Quaternion{T}}, x::Real) where {T} =
    Quaternion(convert(T, x), convert(T, 0), convert(T, 0), convert(T, 0))
convert(::Type{Quaternion{T}}, z::Complex) where {T} =
    Quaternion(convert(T, real(z)), convert(T, imag(z)), convert(T, 0), convert(T, 0))
convert(::Type{Quaternion{T}}, q::Quaternion{T}) where {T <: Real} = q
convert(::Type{Quaternion{T}}, q::Quaternion) where {T} =
    Quaternion(convert(T, q.s), convert(T, q.v1), convert(T, q.v2), convert(T, q.v3), q.norm)

promote_rule(::Type{Quaternion{T}}, ::Type{T}) where {T <: Real} = Quaternion{T}
promote_rule(::Type{Quaternion}, ::Type{T}) where {T <: Real} = Quaternion
promote_rule(::Type{Quaternion{T}}, ::Type{S}) where {T <: Real, S <: Real} = Quaternion{promote_type(T, S)}
promote_rule(::Type{Complex{T}}, ::Type{Quaternion{S}}) where {T <: Real, S <: Real} = Quaternion{promote_type(T, S)}
promote_rule(::Type{Quaternion{T}}, ::Type{Quaternion{S}}) where {T <: Real, S <: Real} = Quaternion{promote_type(T, S)}

quat(p, v1, v2, v3, n = false) = Quaternion(p, v1, v2, v3, n)
quat(x) = Quaternion(x)
quat(s, a) = Quaternion(s, a)

function show(io::IO, q::Quaternion)
    pm(x) = x < 0 ? " - $(-x)" : " + $x"
    print(io, q.s, pm(q.v1), "im", pm(q.v2), "jm", pm(q.v3), "km")
end

real(::Type{Quaternion{T}}) where {T} = T
real(q::Quaternion) = q.s
imag(q::Quaternion) = [q.v1, q.v2, q.v3]

(/)(q::Quaternion, x::Real) = Quaternion(q.s / x, q.v1 / x, q.v2 / x, q.v3 / x)

conj(q::Quaternion) = Quaternion(q.s, -q.v1, -q.v2, -q.v3, q.norm)
abs(q::Quaternion) = sqrt(q.s * q.s + q.v1 * q.v1 + q.v2 * q.v2 + q.v3 * q.v3)
float(q::Quaternion{T}) where T = convert(Quaternion{float(T)}, q)
abs_imag(q::Quaternion) = sqrt(q.v1 * q.v1 + q.v2 * q.v2 + q.v3 * q.v3)
abs2(q::Quaternion) = q.s * q.s + q.v1 * q.v1 + q.v2 * q.v2 + q.v3 * q.v3
inv(q::Quaternion) = q.norm ? conj(q) : conj(q) / abs2(q)

isfinite(q::Quaternion) = q.norm ? true : (isfinite(q.s) && isfinite(q.v1) && isfinite(q.v2) && isfinite(q.v3))

function normalize(q::Quaternion)
    if (q.norm)
        return q
    end
    q = q / abs(q)
    Quaternion(q.s, q.v1, q.v2, q.v3, true)
end

function normalizea(q::Quaternion)
    if (q.norm)
        return (q, one(q.s))
    end
    a = abs(q)
    q = q / a
    (Quaternion(q.s, q.v1, q.v2, q.v3, true), a)
end

function normalizeq(q::Quaternion)
    a = abs(q)
    if a > 0
        q = q / a
        Quaternion(q.s, q.v1, q.v2, q.v3, true)
    else
        Quaternion(0.0, 1.0, 0.0, 0.0, true)
    end
end

(-)(q::Quaternion) = Quaternion(-q.s, -q.v1, -q.v2, -q.v3, q.norm)

(+)(q::Quaternion, w::Quaternion) =
    Quaternion(q.s + w.s, q.v1 + w.v1, q.v2 + w.v2, q.v3 + w.v3)

(-)(q::Quaternion, w::Quaternion) =
    Quaternion(q.s - w.s, q.v1 - w.v1, q.v2 - w.v2, q.v3 - w.v3)

(*)(q::Quaternion, w::Quaternion) = Quaternion(q.s * w.s - q.v1 * w.v1 - q.v2 * w.v2 - q.v3 * w.v3,
                                               q.s * w.v1 + q.v1 * w.s + q.v2 * w.v3 - q.v3 * w.v2,
                                               q.s * w.v2 - q.v1 * w.v3 + q.v2 * w.s + q.v3 * w.v1,
                                               q.s * w.v3 + q.v1 * w.v2 - q.v2 * w.v1 + q.v3 * w.s,
                                               q.norm && w.norm)
(/)(q::Quaternion, w::Quaternion) = q * inv(w)

angleaxis(q::Quaternion) = angle(q), axis(q)

angle(q::Quaternion) = 2 * atan(√(q.v1^2 + q.v2^2 + q.v3^2), q.s)

function axis(q::Quaternion)
    q = normalize(q)
    s = sin(angle(q) / 2)
    abs(s) > 0 ?
        [q.v1, q.v2, q.v3] / s :
        [1.0, 0.0, 0.0]
end

argq(q::Quaternion) = normalizeq(Quaternion(0, q.v1, q.v2, q.v3))

function exp(q::Quaternion)
    s = q.s
    se = exp(s)
    scale = se
    th = abs_imag(q)
    if th > 0
        scale *= sin(th) / th
    end
    Quaternion(se * cos(th), scale * q.v1, scale * q.v2, scale * q.v3, abs(s) < eps(typeof(s)))
end

function log(q::Quaternion)
    q, a = normalizea(q)
    s = q.s
    M = abs_imag(q)
    th = atan(M, s)
    if M > 0
        M = th / M
        return Quaternion(log(a), q.v1 * M, q.v2 * M, q.v3 * M)
    else
        return Quaternion(log(a), th, 0.0, 0.0)
    end
end

function sin(q::Quaternion)
    L = argq(q)
    return (exp(L * q) - exp(-L * q)) / (2 * L)
end

function cos(q::Quaternion)
    L = argq(q)
    return (exp(L * q) + exp(-L * q)) / 2
end

(^)(q::Quaternion, w::Quaternion) = exp(w * log(q))

sqrt(q::Quaternion) = exp(0.5 * log(q))

function linpol(p::Quaternion, q::Quaternion, t::Real)
    p = normalize(p)
    q = normalize(q)
    qm = -q
    if abs(p - q) > abs(p - qm)
        q = qm
    end
    c = p.s * q.s + p.v1 * q.v1 + p.v2 * q.v2 + p.v3 * q.v3
    if c > - 1.0
        if c < 1.0
            o = acos(c)
            s = sin(o)
            sp = sin((1 - t) * o) / s
            sq = sin(t * o) / s
        else
            sp = 1 - t
            sq = t
        end
        Quaternion(sp * p.s  + sq * q.s,
                   sp * p.v1 + sq * q.v1,
                   sp * p.v2 + sq * q.v2,
                   sp * p.v3 + sq * q.v3, true)
    else
        s  =  p.v3
        v1 = -p.v2
        v2 =  p.v1
        v3 = -p.s
        sp = sin((0.5 - t) * pi)
        sq = sin(t * pi)
        Quaternion(s,
                   sp * p.v1 + sq * v1,
                   sp * p.v2 + sq * v2,
                   sp * p.v3 + sq * v3, true)
    end
end

quatrand()  = quat(randn(), randn(), randn(), randn())
nquatrand() = normalize(quatrand())

rand(r::AbstractRNG, ::SamplerType{Quaternion{T}}) where {T<:Real} = quat(rand(r, T), rand(r, T),rand(r,T),rand(r,T))

## Rotations

function qrotation(axis::Vector{T}, theta) where {T <: Real}
    if length(axis) != 3
        error("Must be a 3-vector")
    end
    u = normalize(axis)
    thetaT = convert(eltype(u), theta)
    s = sin(thetaT / 2)
    Quaternion(cos(thetaT / 2), s * u[1], s * u[2], s * u[3], true)
end

# Variant of the above where norm(rotvec) encodes theta
function qrotation(rotvec::Vector{T}) where {T <: Real}
    if length(rotvec) != 3
        error("Must be a 3-vector")
    end
    theta = norm(rotvec)
    if theta > 0
        s = sin(theta / 2) / theta  # divide by theta to make rotvec a unit vector
        return Quaternion(cos(theta / 2), s * rotvec[1], s * rotvec[2], s * rotvec[3], true)
    end
    Quaternion(one(T), zero(T), zero(T), zero(T), true)
end

function qrotation(dcm::Matrix{T}) where {T<:Real}
    # See https://arxiv.org/pdf/math/0701759.pdf
    a2 = 1 + dcm[1,1] + dcm[2,2] + dcm[3,3]
    b2 = 1 + dcm[1,1] - dcm[2,2] - dcm[3,3]
    c2 = 1 - dcm[1,1] + dcm[2,2] - dcm[3,3]
    d2 = 1 - dcm[1,1] - dcm[2,2] + dcm[3,3]

    if a2 >= max(b2, c2, d2)
        a = sqrt(a2)/2
        return Quaternion(a, (dcm[3,2]-dcm[2,3])/4a, (dcm[1,3]-dcm[3,1])/4a, (dcm[2,1]-dcm[1,2])/4a)
    elseif b2 >= max(c2, d2)
        b = sqrt(b2)/2
        return Quaternion((dcm[3,2]-dcm[2,3])/4b, b, (dcm[2,1]+dcm[1,2])/4b, (dcm[1,3]+dcm[3,1])/4b)
    elseif c2 >= d2
        c = sqrt(c2)/2
        return Quaternion((dcm[1,3]-dcm[3,1])/4c, (dcm[2,1]+dcm[1,2])/4c, c, (dcm[3,2]+dcm[2,3])/4c)
    else
        d = sqrt(d2)/2
        return Quaternion((dcm[2,1]-dcm[1,2])/4d, (dcm[1,3]+dcm[3,1])/4d, (dcm[3,2]+dcm[2,3])/4d, d)
    end
end

function qrotation(dcm::Matrix{T}, qa::Quaternion) where {T<:Real}
    q = qrotation(dcm)
    abs(q-qa) < abs(q+qa) ? q : -q
end

rotationmatrix(q::Quaternion) = rotationmatrix_normalized(normalize(q))

function rotationmatrix_normalized(q::Quaternion)
    sx, sy, sz = 2q.s * q.v1, 2q.s * q.v2, 2q.s * q.v3
    xx, xy, xz = 2q.v1^2, 2q.v1 * q.v2, 2q.v1 * q.v3
    yy, yz, zz = 2q.v2^2, 2q.v2 * q.v3, 2q.v3^2
    [1 - (yy + zz)     xy - sz     xz + sy;
        xy + sz   1 - (xx + zz)    yz - sx;
        xz - sy      yz + sx  1 - (xx + yy)]
end

function normalize(v::Vector{T}) where {T}
    nv = norm(v)
    if nv > 0
        return v / nv
    end
    zeros(promote_type(T, typeof(nv)), length(v))
end

function slerp(qa::Quaternion{T}, qb::Quaternion{T}, t::T) where {T}
    # http://www.euclideanspace.com/maths/algebra/realNormedAlgebra/quaternions/slerp/
    coshalftheta = qa.s * qb.s + qa.v1 * qb.v1 + qa.v2 * qb.v2 + qa.v3 * qb.v3;

    if coshalftheta < 0
        qm = -qb
        coshalftheta = -coshalftheta
    else
        qm = qb
    end
    abs(coshalftheta) >= 1.0 && return qa

    halftheta    = acos(coshalftheta)
    sinhalftheta = sqrt(one(T) - coshalftheta * coshalftheta)

    if abs(sinhalftheta) < 0.001
        return Quaternion(
            T(0.5) * (qa.s  + qb.s),
            T(0.5) * (qa.v1 + qb.v1),
            T(0.5) * (qa.v2 + qb.v2),
            T(0.5) * (qa.v3 + qb.v3),
        )
    end

    ratio_a = sin((one(T) - t) * halftheta) / sinhalftheta
    ratio_b = sin(t * halftheta) / sinhalftheta

    Quaternion(
        qa.s  * ratio_a + qm.s  * ratio_b,
        qa.v1 * ratio_a + qm.v1 * ratio_b,
        qa.v2 * ratio_a + qm.v2 * ratio_b,
        qa.v3 * ratio_a + qm.v3 * ratio_b,
    )
end


using ReachabilityAnalysis, JLD2
using ReachabilityAnalysis: add_dimension

LazySets.set_ztol(Float64, 1e-15);

examples_dir = normpath(@__DIR__, "..", "..", "..", "examples")
ISS_path = joinpath(examples_dir, "ISS", "ISS.jld2")

@load ISS_path C;
const C3 = C[3, :]; # variable y₃
const C3_ext = vcat(C3, fill(0.0, 3));

function ISSF01()
    @load ISS_path A B

    U = Hyperrectangle(low=[0.0, 0.8, 0.9], high=[0.1, 1., 1.])
    X0 = BallInf(zeros(size(A, 1)), 0.0001)
    return @ivp(x' = A*x + B*u, x(0) ∈ X0, u ∈ U, x ∈ Universe(270))
end

dirs = CustomDirections([C3, -C3]);
prob_ISSF01 = ISSF01();
sol_ISSF01 = solve(prob_ISSF01, T=20.0, alg=LGG09(δ=6e-4, template=dirs, sparse=true, cache=false));

dim(sol_ISSF01)

πsol_ISSF01 = flatten(sol_ISSF01);

dim(πsol_ISSF01)

using Plots, Plots.PlotMeasures, LaTeXStrings

fig = Plots.plot();
Plots.plot!(fig, πsol_ISSF01[1:10:end], vars=(0, 1), linecolor=:blue, color=:blue, alpha=0.8,
     xlab=L"t",
     ylab=L"y_{3}",
     xtick=[0, 5, 10, 15, 20.], ytick=[-0.00075, -0.0005, -0.00025, 0, 0.00025, 0.0005, 0.00075],
     xlims=(0., 20.), ylims=(-0.00075, 0.00075),
     bottom_margin=6mm, left_margin=2mm, right_margin=4mm, top_margin=3mm);
fig

function ISSC01()
    @load ISS_path A B

    A_ext = add_dimension(A, 3)
    A_ext[1:270, 271:273] = B

    U = Hyperrectangle(low=[0.0, 0.8, 0.9], high=[0.1, 1., 1.])
    X0 = BallInf(zeros(size(A, 1)), 0.0001)
    X0 = X0 * U
    return @ivp(x' = A_ext*x, x(0) ∈ X0)
end

dirs = CustomDirections([C3_ext, -C3_ext]);
prob_ISSC01 = ISSC01();
sol_ISSC01 = solve(prob_ISSC01, T=20.0, alg=LGG09(δ=0.01, template=dirs, sparse=true, cache=false));

πsol_ISSC01 = flatten(sol_ISSC01);

fig = Plots.plot();
Plots.plot!(fig, πsol_ISSC01, vars=(0, 1), linecolor=:blue, color=:blue, alpha=0.8, lw=1.0,
    xlab=L"t",
    ylab=L"y_{3}",
    xtick=[0, 5, 10, 15, 20.], ytick=[-0.0002, -0.0001, 0.0, 0.0001, 0.0002],
    xlims=(0., 20.), ylims=(-0.0002, 0.0002),
    bottom_margin=6mm, left_margin=2mm, right_margin=4mm, top_margin=3mm)
fig


# load init commands
using InteractiveControlPlots, ControlSystems, GtkInteract

# create sliders etc
x = linspace(-1,1,2)
ws = logspace(-2,2,1000)
k=slider(1:0.1:3, label="K")
logT = slider(-1:0.01:1, label="T")
f1 = checkbox(false,label="Bode plot  ")
f2 = checkbox(false,label="Nyquist plot  ")
f3 = checkbox(false,label="PZ map  ")
f4 = checkbox(false,label="Step plot  ")

# create window
cg1 = cairographic(width=200,height=150)
cg2 = cairographic(width=200,height=150)
cg3 = cairographic(width=200,height=150)
cg4 = cairographic(width=200,height=150)
lT = label("T: ")

vbl(w, len) = vbox(label(w.label),hbox(hskip(len),w))
vbl(w) = vbox(label(w.label),w)

w = window(vbox(hbox(vbl(k),
    vbox(lT,logT),
    vbox(vbl(f1,25),vbl(f3,15)),
    vbox(vbl(f2,35),vbl(f4,22))),
    hbox(grow(cg1), grow(cg2)),
    hbox(grow(cg3), grow(cg4))),
    title="Example")

# loop over parameters
Reactive.preserve(map(k,logT,f1,f2,f3,f4) do k,logT,f1,f2,f3,f4
  T = 10^logT

  #Create the 2nd order system
  sys = tf(k*[1,1],[1, T, 1])

  #Create the plots
  fig1 = bode_immerse(sys, ws, 1e-2, 1e2) #Ylims = (1e-2, 1e2)
  fig2 = nyquist_immerse(sys, ws, -2, 2, -2, 2) #x and ylims
  fig3 = pzmap_immerse(sys,-4,4,-4,4) # x and ylims
  y,t,x = step(sys, 10)
  fig4 = timeplot_immerse(t, y, 0, 4) #ylims

  #Update the text
  push!(lT, "T: " * string(round(T,2)))

  #Push figures
  push!(cg1, f1 ? fig1 : nothing)
  push!(cg2, f2 ? fig2 : nothing)
  push!(cg3, f3 ? fig3 : nothing)
  push!(cg4, f4 ? fig4 : nothing)

end)

# output window for display
w


function test_sqg_problemtype(dev, T)
  prob = SurfaceQG.Problem(dev; T=T)

  A = ArrayType(dev)

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

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

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

  prob = SurfaceQG.Problem(dev; calcF=calcF!, stochastic=true)

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

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

Tests the advection term in the SurfaceQG 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 bf (with associated
streamfunction ψf) for which the advection term J(ψf, bf) is non-zero. Next, a
forcing Ff is derived according to Ff = ∂bf/∂t + J(ψf, bf) - ν∇²bf. One solution
to the buoyancy equation forced by this Ff is then bf. (This solution may not
be realized, at least at long times, if it is unstable.)
"""
function test_sqg_advection(dt, stepper, dev::Device=CPU(); n=128, L=2π, ν=1e-2, nν=1)
  n, L  = 128, 2π
  ν, nν = 1e-2, 1
  tf = 0.5  # SQG piles up energy at small scales so running for t ⪆ 0.5 brings instability
  nt = round(Int, tf/dt)

  grid = TwoDGrid(dev, n, L)
  x, y = gridpoints(grid)

  ψf = @.             sin(2x)*cos(2y) +            2sin(x)*cos(3y)
  bf = @. - sqrt(8) * sin(2x)*cos(2y) - sqrt(10) * 2sin(x)*cos(3y)

  Ff = @. (
    - ν*(16*sqrt(2) * sin(2x)*cos(2y) + 10*sqrt(10) * 2sin(x)*cos(3y) )
    - 4*sqrt(2)*(sqrt(5)-2) * (cos(x)*cos(3y)*sin(2x)*sin(2y)
      - 3cos(2x)*cos(2y)*sin(x)*sin(3y))
  )

  Ffh = rfft(Ff)

  # Forcing
  function calcF!(Fh, sol, t, clock, vars, params, grid)
    Fh .= Ffh
    return nothing
  end

  prob = SurfaceQG.Problem(dev; nx=n, Lx=L, ν=ν, nν=nν, dt=dt, stepper=stepper, calcF=calcF!, stochastic=false)

  SurfaceQG.set_b!(prob, bf)

  stepforward!(prob, nt)

  SurfaceQG.updatevars!(prob)

  isapprox(prob.vars.b, bf, rtol=rtol_surfaceqg)
end

function test_sqg_kineticenergy_buoyancyvariance(dev::Device=CPU())
  nx, Lx  = 128, 2π
  ny, Ly  = 128, 3π

  grid = TwoDGrid(dev, nx, Lx, ny, Ly)
  x, y = gridpoints(grid)

  k₀, l₀ = 2π/grid.Lx, 2π/grid.Ly # fundamental wavenumbers
  ψ₀ = @. sin(2k₀*x)*cos(2l₀*y) + 2sin(k₀*x)*cos(3l₀*y)
  b₀ = @. - sqrt(8) * sin(2k₀*x)*cos(2l₀*y) - sqrt(10) * 2sin(k₀*x)*cos(3l₀*y)

  kinetic_energy_calc = 6
  buoyancy_variance_calc = 12

  prob = SurfaceQG.Problem(dev; nx=nx, Lx=Lx, ny=ny, Ly=Ly, stepper="ForwardEuler")

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

  SurfaceQG.set_b!(prob, b₀)
  SurfaceQG.updatevars!(prob)

  kinetic_energy_b₀ = SurfaceQG.kinetic_energy(prob)
  buoyancy_variance_b₀ = SurfaceQG.buoyancy_variance(prob)

  params = SurfaceQG.Params(p.ν, p.nν)

  return (isapprox(kinetic_energy_b₀, kinetic_energy_calc, rtol=rtol_surfaceqg) &&
   isapprox(buoyancy_variance_b₀, buoyancy_variance_calc, rtol=rtol_surfaceqg))
end

function test_sqg_paramsconstructor(dev::Device=CPU())
  n, L = 128, 2π
  ν, nν = 1e-3, 4

  prob = SurfaceQG.Problem(dev; nx=n, Lx=L, ν=ν, nν=nν, stepper="ForwardEuler")

  params = SurfaceQG.Params(ν, nν)

  return prob.params == params
end

function test_sqg_noforcing(dev::Device=CPU())
  n, L = 16, 2π

  prob_unforced = SurfaceQG.Problem(dev; nx=n, Lx=L, stepper="ForwardEuler")

  SurfaceQG.addforcing!(prob_unforced.timestepper.N, prob_unforced.sol, prob_unforced.clock.t, prob_unforced.clock, prob_unforced.vars, prob_unforced.params, prob_unforced.grid)

  function calcF!(Fh, sol, t, clock, vars, params, grid)
    Fh .= 2 * ArrayType(dev)(ones(size(sol)))
    return nothing
  end

  prob_forced = SurfaceQG.Problem(dev; nx=n, Lx=L, stepper="ForwardEuler", calcF=calcF!)

  SurfaceQG.addforcing!(prob_forced.timestepper.N, prob_forced.sol, prob_forced.clock.t, prob_forced.clock, prob_forced.vars, prob_forced.params, prob_forced.grid)

  return prob_unforced.timestepper.N == Complex.(ArrayType(dev)(zeros(size(prob_unforced.sol)))) && prob_forced.timestepper.N == Complex.(2*ArrayType(dev)(ones(size(prob_unforced.sol))))
end

function test_sqg_deterministicforcing_buoyancy_variance_budget(dev::Device=CPU(); n=256, dt=0.01, L=2π, ν=1e-7, nν=2, tf=10.0)
  n, L  = 256, 2π
  ν, nν = 1e-7, 2
  dt, tf = 0.005, 10
  nt = round(Int, tf/dt)

  grid  = TwoDGrid(dev, n, L)
  x, y = gridpoints(grid)

  # buoyancy forcing = 0.01cos(4x)cos(5y)cos(2t)
  f = @. 0.01cos(4x) * cos(5y)
  fh = rfft(f)
  function calcF!(Fh, sol, t, clock, vars, params, grid)
    @. Fh = fh * cos(2t)
    return nothing
  end

  prob = SurfaceQG.Problem(dev; nx=n, Lx=L, ν=ν, nν=nν, dt=dt, stepper="RK4",
    calcF=calcF!, stochastic=false)

  SurfaceQG.set_b!(prob, 0*x)

  B = Diagnostic(SurfaceQG.buoyancy_variance,    prob, nsteps=nt)
  D = Diagnostic(SurfaceQG.buoyancy_dissipation, prob, nsteps=nt)
  W = Diagnostic(SurfaceQG.buoyancy_work,        prob, nsteps=nt)
  diags = [B, D, W]

  stepforward!(prob, diags, nt)

  SurfaceQG.updatevars!(prob)

  dBdt_numerical = (B[3:B.i] - B[1:B.i-2]) / (2 * prob.clock.dt)
  dBdt_computed  = W[2:B.i-1] - D[2:B.i-1]

  return isapprox(dBdt_numerical, dBdt_computed, rtol = 1e-4)
end

function test_sqg_stochasticforcing_buoyancy_variance_budget(dev::Device=CPU(); n=256, L=2π, dt=0.005, ν=1e-7, nν=2, tf=2.0)
  nt = round(Int, tf/dt)

  # Forcing parameters
  kf, dkf = 12.0, 2.0
  εᵇ = 0.01

  grid = TwoDGrid(dev, n, L)
  x, y = gridpoints(grid)

  Kr = ArrayType(dev)([CUDA.@allowscalar grid.kr[i] for i=1:grid.nkr, j=1:grid.nl])

  forcing_spectrum = ArrayType(dev)(zero(grid.Krsq))
  @. forcing_spectrum = exp(-(sqrt(grid.Krsq) - kf)^2 / (2 * dkf^2))
  @. forcing_spectrum = ifelse(grid.Krsq < 2^2, 0, forcing_spectrum)
  @. forcing_spectrum = ifelse(grid.Krsq > 20^2, 0, forcing_spectrum)
  @. forcing_spectrum = ifelse(Kr < 2π/L, 0, forcing_spectrum)
  εᵇ0 = parsevalsum(forcing_spectrum, grid) / (grid.Lx * grid.Ly)
  forcing_spectrum .= εᵇ / εᵇ0 * forcing_spectrum

  Random.seed!(1234)

  function calcF!(Fh, sol, t, clock, vars, params, grid)
    eta = ArrayType(dev)(exp.(2π * im * rand(Float64, size(sol))) / sqrt(clock.dt))
    CUDA.@allowscalar eta[1, 1] = 0.0
    @. Fh = eta * sqrt(forcing_spectrum)
    nothing
  end

  prob = SurfaceQG.Problem(dev; nx=n, Lx=L, ν=ν, nν=nν, dt=dt, stepper="RK4",
    calcF=calcF!, stochastic=true)

  SurfaceQG.set_b!(prob, 0*x)

  B = Diagnostic(SurfaceQG.buoyancy_variance,    prob, nsteps=nt)
  D = Diagnostic(SurfaceQG.buoyancy_dissipation, prob, nsteps=nt)
  W = Diagnostic(SurfaceQG.buoyancy_work,        prob, nsteps=nt)
  diags = [B, D, W]

  stepforward!(prob, diags, nt)

  SurfaceQG.updatevars!(prob)

  dBdt_numerical = (B[2:B.i] - B[1:B.i-1]) / prob.clock.dt

  dBdt_computed = W[2:B.i] - D[1:B.i-1]

  return isapprox(dBdt_numerical, dBdt_computed, rtol = 1e-4)
end


using Test
using PETSc, MPI, LinearAlgebra, SparseArrays

MPI.Initialized() || MPI.Init()

@testset "SNES" begin

  comm    = MPI.COMM_WORLD
  mpirank = MPI.Comm_rank(comm)
  mpisize = MPI.Comm_size(comm)

  for petsclib in PETSc.petsclibs
      PETSc.initialize(petsclib)
      PetscScalar = PETSc.scalartype(petsclib)
      PetscInt    = PETSc.inttype(petsclib)

      # This shows some examples of how SNES can be used.
      # Note that you will receive pointers to PETSc vectors within your user routines.
      #       That is important for parallel simulations, where global residual
      #       and solution vectors are passed to the user-routines, but we work with
      #       local vectors

      function fn!(cfx, cx, user_ctx)
        # We could do Global->Local here on cfx/cx, provided a pointer to the local
        #  vector is available in user_ctx
        x_in  = PETSc.unsafe_localarray(PetscScalar, cx;  write=false)   # read array
        fx_in = PETSc.unsafe_localarray(PetscScalar, cfx; read=false)    # write array

        fx_in[1] = x_in[1]^2       +  x_in[1]*x_in[2] - 3
        fx_in[2] = x_in[1]*x_in[2] +  x_in[2]^2 - 6

        Base.finalize(fx_in)
        Base.finalize(x_in)
      end

      function update_jac!(cx, J1, args...)
          x_in  = PETSc.unsafe_localarray(PetscScalar, cx;  write=false)
          J1[1,1] = 2x_in[1] + x_in[2]
          J1[1,2] = x_in[1]
          J1[2,1] = x_in[2]
          J1[2,2] = x_in[1] + 2x_in[2]

          Base.finalize(x_in)
          PETSc.assemble(J1)
      end

      # structure with which we can pass data to the user-routines above
      mutable struct Data
          vec
          julia
      end

      julia_vec = 0;  # we want pointers to local vectors

      # NOTE: for some reason, snes_monitor and ksp_monitor is not working if we specify that here.
      # To be sorted out why that is (converged_reason does work)
      S = PETSc.SNES{PetscScalar}(petsclib,comm;
                      ksp_rtol=1e-4,
                      pc_type="none",
                      ksp_monitor=true,
                      snes_monitor=true,
                      snes_converged_reason=false,
                      ksp_converged_reason=false)

      data        = Data(PetscScalar.([100;2]), 1)
      S.user_ctx  = data;      # we can pack anything we need in this struct

      PJ = PETSc.MatSeqDense(zeros(PetscScalar,(2,2)))
      PETSc.setfunction!(S, fn!, PETSc.VecSeq( PetscScalar.(zeros(2))))
      PETSc.setjacobian!(S, update_jac!, PJ, PJ)

      x  = PETSc.VecSeq(PetscScalar.([2.0, 3.0]));
      b  = PETSc.VecSeq(PetscScalar.([0.0, 0.0]));
      PETSc.solve!(x, S, b)

      sol = PETSc.unsafe_localarray(PetscScalar, x.ptr; read=false)
      @test sol ≈ [1.0,2.0] rtol=1e-4

      # cleanup
      PETSc.destroy(x);
      PETSc.destroy(b);
      PETSc.destroy(PJ);
      #PETSc.destroy(S);

      PETSc.finalize(petsclib)
    end
end


using Revise
using Distributed, MPIClusterManagers

manager = MPIManager(np=4)
addprocs(manager)

@everywhere ENV["OMP_NUM_THREADS"] = 4

@everywhere using Revise
@everywhere using Devito
@everywhere using MPI
@everywhere configuration!("log-level", "DEBUG")
@everywhere configuration!("language", "openmp")
@everywhere configuration!("mpi", true)

@everywhere function ricker(f, _t, t₀)
    t = reshape(_t, length(_t), 1)
    (1.0 .- 2.0 * (pi * f * (t .- t₀)).^2) .* exp.(-(pi * f * (t .- t₀)).^2)
end

@everywhere function model()
    write(stdout, "inside model()\n")
    grid = Grid(
        shape = (121,121,121),
        origin = (0.0,0.0,0.0),
        extent = (1210.0,1210.0,1210.0),
        dtype = Float32)

    b = Devito.Function(name="b", grid=grid, space_order=8)
    v = Devito.Function(name="vel", grid=grid, space_order=8)
    q = Devito.Function(name="woverq", grid=grid, space_order=8)

    b_data = data(b)
    v_data = data(v)
    q_data = data(q)

    b_data .= 1
    v_data .= 1.5
    q_data .= 1/1000

    time_range = 0.0f0:1.0f0:750.0f0

    p = TimeFunction(name="p", grid=grid, time_order=2, space_order=8)
    z,y,x,t = dimensions(p)

    src = SparseTimeFunction(name="src", grid=grid, f0=0.01f0, npoint=1, nt=length(time_range), coordinates=[605.0 605.0 10.0])
    src_data = data(src)
    w = ricker(0.01, collect(time_range), 125)
    copy!(src_data, w)
    src_term = inject(src; field=forward(p), expr=src * spacing(t)^2 * v^2 / b)

    nz,ny,nx,δz,δy,δx = size(grid)...,spacing(grid)...
    rec = SparseTimeFunction(name="rec", grid=grid, npoint=nx, nt=length(time_range))
    rec_coords = coordinates(rec)
    _rec_coords = zeros(Float32, length(dimensions(p))-1, nx)
    _rec_coords[1,:] .= δx*(0:nx-1)
    _rec_coords[2,:] .= 605
    _rec_coords[3,:] .= 20
    copy!(rec_coords, _rec_coords)
    rec_term = interpolate(rec, expr=p)

    g1(field) = dx(field,x0=x+spacing(x)/2)
    g2(field) = dy(field,x0=y+spacing(y)/2)
    g3(field) = dz(field,x0=z+spacing(z)/2)
    g1_tilde(field) = dx(field,x0=x-spacing(x)/2)
    g2_tilde(field) = dy(field,x0=y-spacing(y)/2)
    g3_tilde(field) = dz(field,x0=z-spacing(z)/2)

    # write the time update equation for u_x
    update_p = spacing(t)^2 * v^2 / b *
        (g1_tilde(b * g1(p)) + g2_tilde(b * g2(p)) + g3_tilde(b * g3(p))) +
        (2 - spacing(t) * q) * p +
        (spacing(t) * q - 1) * backward(p)

    stencil_p = Eq(forward(p), update_p)

    # update the dimension spacing_map to include the time dimension
    # these symbols will be replaced with the relevant scalars by the Operator
    dt = step(time_range)
    smap = spacing_map(grid)
    smap[spacing(t)] = dt

    op = Operator([stencil_p, src_term, rec_term], subs=smap, name="OpExampleIso")

    bx,by = 19,8
    t_apply = @elapsed begin
        summary = apply(op; x0_blk0_size=bx, y0_blk0_size=by)
    end

    write(stdout, "t_appy=$t_apply\n")

    _p = data_with_halo(p)
    _d = data(rec)

    __d = convert(Array, _d)

    if MPI.Comm_rank(MPI.COMM_WORLD) == 0
        write("d.bin", __d)
    end
    MPI.Barrier(MPI.COMM_WORLD)

    nothing
end

@mpi_do manager model()

d = read!("d.bin", Array{Float32}(undef,751,121))

using PyPlot
figure(); imshow(d); display(gcf())


function benchmark(bd::BenchmarkData, method::RegressionMethod, nfolds::Int=5)

    nrange = length(bd.n)
    results = zeros(nrange, 6)

    validation_results = Array{ValidationResults,2}(nrange, nfolds)

    for i = 1:nrange
        # Generate some synthetic data
        rd = getdata(bd, i)
        true_support = find(abs.(rd.w) .> 1e-6)

        folds = kfolds(shuffleobs((rd.X', rd.Y)), k = nfolds)
        j = 0
        for ((X_train, y_train), (X_test, y_test)) in folds
            j += 1
            # Validate the right method parameters
            X_train, X_test = X_train', X_test'
            v_results = validate_params!(X_train, y_train, bd.sparsity, method)
            validation_results[i, j] = v_results
            tic()
            indices, w = solve_problem(method, X_train, y_train, bd.sparsity)
            time = toq()
            a = accuracy(indices, true_support)
            f = falsepositive(indices, true_support)
            pred_train = predict_sparse(X_train, indices, w)
            train_R2 = isRsquared(pred_train, y_train)
            pred_test  = predict_sparse(X_test, indices, w)
            test_R2 = oosRsquared(pred_test, y_test, y_train)
            @show bd.n[i], a
            results[i, :] .+= [a, f, train_R2, test_R2, time, method.gamma]
        end
    end
    results / nfolds, validation_results
end


using LinearAlgebra

mutable struct Broyden{R <: Real, C <: Union{R, Complex{R}}, T <: AbstractArray{C}}
    H
    theta_bar::R
    function Broyden{R, C, T}(x::T, H, theta_bar) where {R, C, T}
        new(H, theta_bar)
    end
end

Broyden(x::T; H=I, theta_bar=R(0.2)) where {
    R <: Real, C <: Union{R, Complex{R}}, T <: AbstractArray{C}
} = Broyden{R, C, T}(x, H, theta_bar)

_sign(x::R) where R = x == 0 ? R(1) : sign(x)

function update!(L::Broyden{R, C, T}, s, y) where {R, C, T}
    Hy = L.H * y
    sH = s' * L.H
    delta = dot(Hy, s)/norm(s)^2
    theta = if abs(delta) >= L.theta_bar
        R(1)
    else
        (1 - _sign(delta) * L.theta_bar)/(1 - delta)
    end
    L.H += (s - Hy)/dot(s, (1/theta - 1) * s + Hy) * sH
end

import Base: *

function (*)(L::Broyden, v)
    w = similar(v)
    mul!(w, L, v)
end

import LinearAlgebra: mul!

mul!(d::T, L::Broyden, v::T) where {T} = mul!(d, L.H, v)


####################################################################
# ML estimation
####################################################################

model_g1 = Sem(
    specification = specification_g1,
    data = dat_g1,
    imply = RAMSymbolic
)

model_g2 = Sem(
    specification = specification_g2,
    data = dat_g2,
    imply = RAM
)

model_ml_multigroup = SemEnsemble(model_g1, model_g2; diff = semdiff)

############################################################################
### test gradients
############################################################################

@testset "ml_gradients_multigroup" begin
    @test test_gradient(model_ml_multigroup, start_test; atol = 1e-9)
end

# fit
@testset "ml_solution_multigroup" begin
    solution = sem_fit(model_ml_multigroup)
    update_estimate!(partable, solution)
    @test compare_estimates(
        partable,
        solution_lav[:parameter_estimates_ml]; atol = 1e-4,
        lav_groups = Dict(:Pasteur => 1, :Grant_White => 2))
end

@testset "fitmeasures/se_ml" begin
    solution_ml = sem_fit(model_ml_multigroup)
    @test all(test_fitmeasures(
        fit_measures(solution_ml),
        solution_lav[:fitmeasures_ml]; rtol = 1e-2, atol = 1e-7))

    update_partable!(
        partable, identifier(model_ml_multigroup), se_hessian(solution_ml), :se)
    @test compare_estimates(
        partable,
        solution_lav[:parameter_estimates_ml]; atol = 1e-3,
        col = :se, lav_col = :se,
        lav_groups = Dict(:Pasteur => 1, :Grant_White => 2))
end

####################################################################
# ML estimation - user defined loss function
####################################################################

struct UserSemML <: SemLossFunction
    objective
    gradient
    hessian
end

############################################################################
### constructor
############################################################################

UserSemML(;n_par, kwargs...) = UserSemML([1.0], zeros(n_par), zeros(n_par, n_par))

############################################################################
### functors
############################################################################

import LinearAlgebra: Symmetric, cholesky, isposdef, logdet, tr
import LinearAlgebra

function (semml::UserSemML)(par, F, G, H, model)
    if G error("analytic gradient of ML is not implemented (yet)") end
    if H error("analytic hessian of ML is not implemented (yet)") end

    a = cholesky(Symmetric(model.imply.Σ); check = false)
    if !isposdef(a)
        semml.objective[1] = Inf
    else
        ld = logdet(a)
        Σ_inv = LinearAlgebra.inv(a)
        if !isnothing(F)
            prod = Σ_inv*model.observed.obs_cov
            semml.objective[1] = ld + tr(prod)
        end
    end
end

# models
model_g1 = Sem(
    specification = specification_g1,
    data = dat_g1,
    imply = RAMSymbolic
)

model_g2 = SemFiniteDiff(
    specification = specification_g2,
    data = dat_g2,
    imply = RAMSymbolic,
    loss = UserSemML
)

model_ml_multigroup = SemEnsemble(model_g1, model_g2; diff = semdiff)

@testset "gradients_user_defined_loss" begin
    @test test_gradient(model_ml_multigroup, start_test; atol = 1e-9)
end

# fit
@testset "solution_user_defined_loss" begin
    solution = sem_fit(model_ml_multigroup)
    update_estimate!(partable, solution)
    @test compare_estimates(
        partable,
        solution_lav[:parameter_estimates_ml]; atol = 1e-4,
        lav_groups = Dict(:Pasteur => 1, :Grant_White => 2))
end

####################################################################
# GLS estimation
####################################################################

model_ls_g1 = Sem(
    specification = specification_g1,
    data = dat_g1,
    imply = RAMSymbolic,
    loss = SemWLS
)

model_ls_g2 = Sem(
    specification = specification_g2,
    data = dat_g2,
    imply = RAMSymbolic,
    loss = SemWLS
)

model_ls_multigroup = SemEnsemble(model_ls_g1, model_ls_g2; diff = semdiff)

@testset "ls_gradients_multigroup" begin
    @test test_gradient(model_ls_multigroup, start_test; atol = 1e-9)
end

@testset "ls_solution_multigroup" begin
    solution = sem_fit(model_ls_multigroup)
    update_estimate!(partable, solution)
    @test compare_estimates(
        partable,
        solution_lav[:parameter_estimates_ls]; atol = 1e-4,
        lav_groups = Dict(:Pasteur => 1, :Grant_White => 2))
end

@testset "fitmeasures/se_ls" begin
    solution_ls = sem_fit(model_ls_multigroup)
    @test all(test_fitmeasures(
        fit_measures(solution_ls),
        solution_lav[:fitmeasures_ls];
        fitmeasure_names = fitmeasure_names_ls, rtol = 1e-2, atol = 1e-5))

    update_partable!(
        partable, identifier(model_ls_multigroup), se_hessian(solution_ls), :se)
    @test compare_estimates(
        partable,
        solution_lav[:parameter_estimates_ls]; atol = 1e-2,
        col = :se, lav_col = :se,
        lav_groups = Dict(:Pasteur => 1, :Grant_White => 2))
end

if !isnothing(specification_miss_g1)

####################################################################
# FIML estimation
####################################################################

model_g1 = Sem(
    specification = specification_miss_g1,
    observed = SemObsMissing,
    loss = SemFIML,
    data = dat_miss_g1,
    imply = RAM,
    diff = SemDiffEmpty()
)

model_g2 = Sem(
    specification = specification_miss_g2,
    observed = SemObsMissing,
    loss = SemFIML,
    data = dat_miss_g2,
    imply = RAM,
    diff = SemDiffEmpty()
)

model_ml_multigroup = SemEnsemble(model_g1, model_g2; diff = semdiff)

############################################################################
### test gradients
############################################################################

start_test = [
    fill(0.5, 6);
    fill(1.0, 9);
    0.05; 0.01; 0.01; 0.05; 0.01; 0.05;
    fill(0.01, 9);
    fill(1.0, 9);
    0.05; 0.01; 0.01; 0.05; 0.01; 0.05;
    fill(0.01, 9)]

@testset "fiml_gradients_multigroup" begin
    @test test_gradient(model_ml_multigroup, start_test; atol = 1e-7)
end

@testset "fiml_solution_multigroup" begin
    solution = sem_fit(model_ml_multigroup)
    update_estimate!(partable_miss, solution)
    @test compare_estimates(
        partable_miss,
        solution_lav[:parameter_estimates_fiml]; atol = 1e-4,
        lav_groups = Dict(:Pasteur => 1, :Grant_White => 2))
end

@testset "fitmeasures/se_fiml" begin
    solution = sem_fit(model_ml_multigroup)
    @test all(test_fitmeasures(
        fit_measures(solution),
        solution_lav[:fitmeasures_fiml]; rtol = 1e-3, atol = 0))

    update_partable!(
        partable_miss, identifier(model_ml_multigroup), se_hessian(solution), :se)
    @test compare_estimates(
        partable_miss,
        solution_lav[:parameter_estimates_fiml]; atol = 1e-3,
        col = :se, lav_col = :se,
        lav_groups = Dict(:Pasteur => 1, :Grant_White => 2))
end

end


"""
# Adam
## Constructor
```julia
    Adam(; alpha=0.0001, beta_mean=0.9, beta_var=0.999, epsilon=1e-8)
```
## Description
Adam is a gradient based optimizer that choses its search direction by building up estimates of the first two moments of the gradient vector. This makes it suitable for problems with a stochastic objective and thus gradient. The method is introduced in [1] where the related AdaMax method is also introduced, see `?AdaMax` for more information on that method.

## References
[1] https://arxiv.org/abs/1412.6980
"""
struct Adam{T}
    α::T
    β₁::T
    β₂::T
    ϵ::T
end
Adam(; alpha = 0.0001, beta_mean = 0.9, beta_var = 0.999, epsilon = 1e-8) =
    Adam(alpha, beta_mean, beta_var, epsilon)
summary(::Adam) = "Adam"

"""
    AdaMax(; alpha=0.002, beta_mean=0.9, beta_var=0.999)
# Adam
## Constructor
```julia
    AdaMax(; alpha=0.002, beta_mean=0.9, beta_var=0.999)
```
## Description
AdaMax is a gradient based optimizer that choses its search direction by building up estimates of the first two moments of the gradient vector. This makes it suitable for problems with a stochastic objective and thus gradient. The method is introduced in [1] where the related Adam method is also introduced, see `?Adam` for more information on that method.

[1] https://arxiv.org/abs/1412.6980
"""
struct AdaMax{T}
    α::T
    β₁::T
    β₂::T
end
AdaMax(; alpha = 0.002, beta_mean = 0.9, beta_var = 0.999) =
    AdaMax(alpha, beta_mean, beta_var)
summary(::AdaMax) = "AdaMax"

function solve(problem::OptimizationProblem, x0, adam::Adam, options::OptimizationOptions)
    α, β₁, β₂, ϵ = adam.α, adam.β₁, adam.β₂, adam.ϵ
    t0 = time()

    fz, ∇fz = upto_gradient(problem, copy(x0), x0)
    f0, ∇f0 = fz, norm(∇fz, Inf)

    z = copy(x0)
    m = copy(∇fz)
    v = fill(zero(∇fz[1]^2), length(m))
    a = 1 - β₁
    b = 1 - β₂

    iter = 0
    is_converged = false

    while !is_converged
        iter = iter + 1
        m = β₁ .* m .+ a .* ∇fz
        v = β₂ .* v .+ b .* ∇fz .^ 2
        # m̂ = m./(1-β₁^t)
        # v̂ = v./(1-β₂^t)
        # x = x .- α*m̂/(sqrt.(v̂+ϵ))
        αₜ = α * sqrt(1 - β₂^iter) / (1 - β₁^iter)
        z = z .- αₜ .* m ./ (sqrt.(v) .+ ϵ)

        # ∇fz = gradient!(mp, x, ∇fz)
        fz, ∇fz = upto_gradient(problem, ∇fz, z)
        is_converged = iter >= options.maxiter
    end
    ConvergenceInfo(
        adam,
        (
            solution = z,
            minimum = fz,
            ∇fz = ∇fz,
            f0 = f0,
            ∇f0 = ∇f0,
            iter = iter,
            time = time() - t0,
        ),
        options,
    )
end
function solve(problem::OptimizationProblem, x0, adam::AdaMax, options::OptimizationOptions)
    α, β₁, β₂ = adam.α, adam.β₁, adam.β₂
    t0 = time()

    fz, ∇fz = upto_gradient(problem, copy(x0), x0)
    f0, ∇f0 = fz, norm(∇fz, Inf)

    z = copy(x0)
    m = copy(∇fz)
    u = fill(zero(∇fz[1]^2), length(m))
    a = 1 - β₁

    iter = 0
    is_converged = false

    while !is_converged
        iter = iter + 1
        m = β₁ .* m .+ a .* ∇fz
        u = max.(β₂ .* u, abs.(∇fz))
        z = z .- (α ./ (1 - β₁^iter)) .* m ./ u

        fz, ∇fz = upto_gradient(problem, ∇fz, z)
        is_converged = iter >= options.maxiter
    end
    ConvergenceInfo(
        adam,
        (
            solution = z,
            minimum = fz,
            ∇fz = ∇fz,
            f0 = f0,
            ∇f0 = ∇f0,
            iter = iter,
            time = time() - t0,
        ),
        options,
    )
end


#=
ncg.jl
Nonlinear CG optimization
2019-03-16, Jeff Fessler, University of Michigan
=#

export ncg

using LinearAlgebra: I, norm

# for tests:
using Test: @test, @test_throws, @inferred
using Plots: Plot
using LinearAlgebra: opnorm
using Random: seed!
using Plots: plot, plot!, scatter!
using LaTeXStrings

"""
`(x,out) = ncg(B, gradf, curvf, x0; ...)`

Nonlinear preconditioned conjugate gradient algorithm
to minimize a general "inverse problem" cost function of the form
``\\Psi(x) = \\sum_{j=1}^J f_j(B_j x)``
where each function ``f_j(v)`` has a quadratic majorizer of the form
`` q_j(v;u) = f_j(u) + \\nabla f_j(u) (v - u) + 1/2 \\|v - \\u|^2_{C_j(u)} ``
where ``C_j(u)`` is diagonal matrix of curvatures.
(It suffices for each ``f_j`` to have a Lipschitz smooth gradient.)

This CG method uses a majorize-minimize (MM) line search.

in
- `B`		array of ``J`` blocks ``B_1,...,B_J``
- `gradf`	array of ``J`` functions return gradients of ``f_1,...,f_J``
- `curvf`	array of ``J`` functions `z -> curv(z)` that return a scalar
		or a vector of curvature values for each element of ``z``
- `x0`	initial guess; need `length(x) == size(B[j],2)` for ``j=1...J``

option
- `niter`	# number of outer iterations; default 50
- `ninner`	# number of inner iterations of MM line search; default 5
- `P`		# preconditioner; default `I`
- `betahow`	"beta" method for the search direction; default `:dai_yuan`
- `fun`		User-defined function to be evaluated with two arguments (x,iter).
			It is evaluated at (x0,0) and then after each iteration.

output
- `x`		final iterate
- `out`		`[niter+1] (fun(x0,0), fun(x1,1), ..., fun(x_niter,niter))`
	(all 0 by default). This is an array of length `niter+1`
"""
function ncg(
		B::AbstractVector{<:Any},
		gradf::AbstractVector{<:Function},
		curvf::AbstractVector{<:Function},
		x0::AbstractVector{<:Number} ;
		niter::Int = 50,
		ninner::Int = 5,
		P=I,
		betahow::Symbol=:dai_yuan,
		fun::Function = (x,iter) -> 0,
	)

	out = Array{Any}(undef, niter+1)
	out[1] = fun(x0, 0)

	J = length(B)

	x = copy(x0)
	dir = []
	grad_old = []
	grad_new = []

	Bx = [B[j] * x for j=1:J] # u_j in course notes
	grad = (Bx) -> sum([B[j]' * gradf[j](Bx[j]) for j=1:J])

for iter = 1:niter
	grad_new = grad(Bx) # gradient
	npgrad = -(P * grad_new)
	if iter == 1
		dir = npgrad
	else
		if betahow === :dai_yuan
			denom =	(grad_new - grad_old)' * dir
			if denom == 0
				betaval = 0
			else
				betaval = grad_new' * (P * grad_new) / denom
			end
		else
			throw(ArgumentError("unknown beta choice: $betahow"))
		end
		dir = npgrad + betaval * dir # search direction
	end
	grad_old = grad_new

	# MM-based line search for step size alpha
	# using h(a) = sum_j f_j(uj + a vj)
	Bd = [B[j] * dir for j=1:J] # v_j in course notes

	alf = 0
	for ii=1:ninner
	#	derh = alf -> sum([Bd[j]' * gradf[j](Bx[j] + alf * Bd[j]) for j=1:J])
		derh = 0 # derivative of h(a)
		curv = 0
		for j=1:J
			tmp = Bx[j] + alf * Bd[j]
			derh += real(Bd[j]' * gradf[j](tmp))
			curv += sum(curvf[j](tmp) .* abs.(Bd[j]).^2)
		end
		curv < 0 && throw("bug: curv < 0")
		if curv > 0
			alf = alf - derh / curv
		end
		if alf == 0
			break
		end
	end

	x += alf * dir
	for j=1:J # update Bj * x
		Bx[j] += alf * Bd[j]
	end
	out[iter+1] = fun(x, iter)
end

	return x, out
#	return eltype(x0).(x), out # todo
end

"""
`(x,out) = ncg(grad, curv, x0, ...)`

special case of `ncg` (nonlinear CG) for minimizing a cost function
whose gradient is `grad(x)`
and that has a quadratic majorizer with diagonal Hessian given by
`curv(x)`.
Typically `curv = (x) -> L` where `L` is the Lipschitz constant of `grad`
"""
function ncg(
		grad::Function,
		curv::Function,
		x0::AbstractVector{<:Number} ;
		kwargs...)

	return ncg([I], [grad], [curv], x0; kwargs...)
end

function ncg_test()
	seed!(0); M = 40; N = 10; A = randn(M,N); y = randn(M)
	a2 = opnorm(A)^2
	reg = 0.1 * a2
	xh = (A'A + reg*I) \ A'y
	cost = (x) -> 1/2 * norm(A * x - y)^2 + reg/2 * norm(x)^2
	fun = (x,iter) -> (cost(x) - cost(xh), norm(x - xh) / norm(xh), time())
	grad1 = (x) -> A' * (A * x - y) + reg * x
	curv1 = (x) -> a2 + reg

	B = [A, I] # matrix blocks
	gradf = [u -> u - y, v -> reg * v] # f functions gradients
	curvf = [v -> 1, v -> reg]

#	todo-i: these next two fail:
#	@inferred ncg([ones(3,3)], [x -> x], [v -> 1], zeros(Float32, 3), niter=2)
#	@inferred ncg(x -> x, v -> 1, zeros(Float32, 3), niter=2)
#	error("todo")

	niter = 40
	x1, out1 = ncg(   grad1, curv1, zeros(N), niter=niter, fun=fun)
	x2, out2 = ncg(B, gradf, curvf, zeros(N), niter=niter, fun=fun)
	!isapprox(x1, xh) && throw("bug: x1 vs xh")
	!isapprox(x2, xh) && throw("bug: x2 vs xh")

	@test_throws ArgumentError ncg(grad1, curv1, zeros(N), betahow=:test)

	lf = x -> log10(max(x,1e-17))
	costk = out -> lf.([out[k][1] for k=1:niter+1])
	errk = out -> lf.([out[k][2] for k=1:niter+1])
	allk = out -> (costk(out), errk(out))
	cost1, err1 = allk(out1)
	cost2, err2 = allk(out2)

	k = 0:niter
	plot(xlabel="k", ylabel=L"\log(\Psi(x_k) - \Psi(\hat{x}))")
	scatter!(k, cost1, color=:blue, label="cost1")
	scatter!(k, cost2, color=:red, marker=:x, label="cost2")
	p1 = plot!()

	plot(xlabel="k", ylabel=L"\log(\|x_k - \hat{x}\|/\|\hat{x}\|)")
	scatter!(k, err1, color=:blue, label="NRMSD1")
	scatter!(k, err2, color=:red, marker=:x, label="NRMSD2")
	p2 = plot!()
	plot(p1, p2)
end

"""
`ncg(:test)`

self test
"""
function ncg(test::Symbol)
	test != :test && throw("test")
	@test ncg_test() isa Plots.Plot
	true
end

#=
todo
f = x -> x + one(eltype(x))
f = x -> eltype(x).(x)

@inferred f(3)
@inferred f(3.0f0)

function g(x0::AbstractVector{<:Number} ; f::Function = x -> 2*x)
	x = copy(x0)
#	return x .+ ones(eltype(x0), size(x0))
#	return 2*x
	return f(x)
end

g(1:3)
@inferred g(1:3)
@code_warntype g(1:3)

#	the following test returns nothing!?
#	@code_warntype ncg([x -> 2x], [v -> 1], zeros(Float32, 3), niter=2)
=#


cmap_tum = ColorScheme(
    reverse(range(
        RGB{Float64}(1,       1,       1),
        RGB{Float64}(0, 101/255, 189/255),
        length=255,
    ))
)

cmap_tum2 = ColorScheme(
    [
        range(
            RGB{Float64}(  0/255,  51/255,  89/255),
            RGB{Float64}(100/255, 160/255, 200/255),
            length=32,
        )
        range(
            RGB{Float64}(100/255, 160/255, 200/255),
            RGB{Float64}(152/255, 198/255, 234/255),
            length=32,
        )
        range(
            RGB{Float64}(152/255, 198/255, 234/255),
            RGB{Float64}(218/255, 215/255, 203/255),
            length=64,
        )
        # range(
        #     RGB{Float64}(218/255, 215/255, 203/255),
        #     RGB{Float64}(162/255, 173/255,   0/255),
        #     length=64,
        # )
        # range(
        #     RGB{Float64}(162/255, 173/255,   0/255),
        #     RGB{Float64}(227/255, 114/255,  34/255),
        #     length=64,
        # )
        range(
            RGB{Float64}(218/255, 215/255, 203/255),
            RGB{Float64}(227/255, 114/255,  34/255),
            length=128,
        )
    ]
)


using NonlinearEigenproblemsTest
using NonlinearEigenproblems
using Test
using LinearAlgebra

@bench @testset "fiber native" begin
    nep = nep_gallery("nlevp_native_fiber")
    n = size(nep, 1)

    # An exact eigenvalue according (reported in NLEVP collection)
    sol_val= 7.139494306065948e-07;

    (λ,v)=quasinewton(nep,λ=7.14e-7,v=ones(n),
                      logger=displaylevel, armijo_factor=0.5,armijo_max=10)

    @test abs(λ-sol_val)<1e-10;

    # check that we "maintain" real arithmetic
    vv=real(v/v[1]);
    (λ1,v)=resinv(Float64,nep,λ=7.14e-7,v=vv,
                      logger=displaylevel)
    @test abs(λ-sol_val)<1e-10;

    @test eltype(v)==Float64

    @testset "Errors thrown" begin
        @test_throws MethodError nep_gallery("nlevp_native_fiber", 15)
        @test_throws MethodError nep_gallery("nlevp_native_fiber", t=15)
    end

end


include("SortSparse.jl")
include("KoLesky.jl")
include("Cholesky.jl")
include("CovFuncs.jl")
include("Utils.jl")
using LinearAlgebra
using Distributions
using Distances
using Random
using Match
using JLD
Random.seed!(123)

N = 10000
d = 2
l = 0.5
σ = exp.(randn(N))
ρList = [1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5, 5.0, 5.5, 6.0]
λ = 1.5

# Creating uniform random data
x = rand(d, N)

# Creating output arrays
NNZ = zeros(length(ρList))
KLTrueNaive = zeros(3, length(ρList))
KLNaiveTrue = zeros(3, length(ρList))
KLTrueApp = zeros(3, length(ρList))
KLAppTrue = zeros(3, length(ρList))
KLTrueLarge = zeros(3, length(ρList))
KLLargeTrue = zeros(3, length(ρList))
KLTrueExact = zeros(3, length(ρList))
KLExactTrue = zeros(3, length(ρList))

σ .= max.(σ, eps(Float64))

covfunc12!(r) = matern12!(r,l)
covfunc32!(r) = matern32!(r,l)
covfunc52!(r) = matern52!(r,l)

for (iOrder, maternOrder) in enumerate(["12", "32", "52"])
  # Make sure to avoid NANs
  @match maternOrder begin
    "12" => (covfunc! = covfunc12!)
    "32" => (covfunc! = covfunc32!)
    "52" => (covfunc! = covfunc52!)
  end

  for (iρ, ρ) in enumerate(ρList)

    # reordering points and forming skeletons
    colptr, rowval, P, revP, distances = sortSparseRev(x, ρ, 1)
    xOrd = x[:, P]
    skeletons = construct_skeletons(colptr, rowval, distances, λ)

    # Creating reference solution
    Θ = covfunc!(pairwise(Euclidean(), xOrd, dims=2)) + Diagonal(σ)

    # Creating naive vecchia approximation
    LNaive = assembleL(skeletons, xOrd, covfunc!, σ)
    ΘNaive = inv(Matrix(LNaive)') * inv(Matrix(LNaive))

    # Creating "smart" vecchia approximation
    L = assembleL(skeletons, xOrd, covfunc!)
    U = squareSparse(L) + spdiagm(0 => 1 ./ σ)
    icholU_high_level!(U)
    noisecov = NoiseCov(L, U, σ)
    ΘExact = exactMatrix(noisecov)
    ΘApp = appMatrix(noisecov)

    # Creating smart vecchia approximation with larger stencil for IC0
    ULarge = triu(L * L' + spdiagm(0 => 1 ./ σ))
    icholU_high_level!(ULarge)
    noisecovLarge = NoiseCov(L, ULarge, σ)
    ΘLarge = appMatrix(noisecovLarge)

    KLTrueNaive[iOrder, iρ] = KL(Θ, ΘNaive)
    KLNaiveTrue[iOrder, iρ] = KL(ΘNaive, Θ)

    KLTrueApp[iOrder, iρ] = KL(Θ, ΘApp)
    KLAppTrue[iOrder, iρ] = KL(ΘApp, Θ)

    KLTrueLarge[iOrder, iρ] = KL(Θ, ΘLarge)
    KLLargeTrue[iOrder, iρ] = KL(ΘLarge, Θ)

    KLTrueExact[iOrder, iρ] = KL(Θ, ΘExact)
    KLExactTrue[iOrder, iρ] = KL(ΘExact, Θ)
  end
end

save("./out/jld/KLVaryRho.jld",
     "ρList", ρList,
     "KLTrueNaive", KLTrueNaive,
     "KLNaiveTrue", KLNaiveTrue,
     "KLTrueApp", KLTrueApp,
     "KLAppTrue", KLAppTrue,
     "KLTrueLarge", KLTrueLarge,
     "KLLargeTrue", KLLargeTrue,
     "KLTrueExact", KLTrueExact,
     "KLExactTrue", KLExactTrue)


import brainflow

brainflow.enable_dev_board_logger()

params = brainflow.BrainFlowInputParams()
board_shim = brainflow.BoardShim(Integer(brainflow.SYNTHETIC_BOARD), params)

brainflow.prepare_session(board_shim)
brainflow.start_stream(board_shim)
sleep(5)
brainflow.stop_stream(board_shim)
data = brainflow.get_current_board_data(32, board_shim)
brainflow.release_session(board_shim)

eeg_channels = brainflow.get_eeg_channels(Integer(brainflow.SYNTHETIC_BOARD))
data_first_channel = data[eeg_channels[1], :]

# returns tuple of wavelet coeffs and lengths
wavelet_data = brainflow.perform_wavelet_transform(data_first_channel, "db4", 2)
restored_wavelet_data = brainflow.perform_inverse_wavelet_transform(wavelet_data, length(data_first_channel), "db4", 2)

fft_data = brainflow.perform_fft(data_first_channel)
restored_fft_data = brainflow.perform_ifft(fft_data)

println("Original Data")
println(data_first_channel)
println("Restored from Wavelet Data")
println(restored_wavelet_data)
println("Restored from FFT Data")
println(restored_fft_data)


@testset "150.evaluate-reverse-polish-notation.jl" begin
    @test eval_rpn(["4", "13", "5", "/", "+"]) == 6
    @test eval_rpn(["2", "1", "+", "3", "*"]) == 9
    @test eval_rpn(["10", "6", "9", "3", "+", "-11", "*", "/", "*", "17", "+", "5", "+"]) == 22
    @test eval_rpn(["2", "1", "-", "3", "*"]) == 3
end


"""
    test_1layerqg_rossbywave(stepper, dt, nsteps, dev; kwargs...)

Evolves a Rossby wave and compares with the analytic solution.
"""
function test_1layerqg_rossbywave(stepper, dt, nsteps, dev::Device=CPU(); deformation_radius=Inf)
  nx = 64
  Lx = 2π
   β = 2.0
   μ = 0.0
   ν = 0.0
   T = Float64

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

  prob = SingleLayerQG.Problem(dev; nx=nx, Lx=Lx, eta=eta, deformation_radius=deformation_radius, β=β, μ=μ, ν=ν, stepper=stepper, dt=dt)
  sol, clock, vars, params, grid = prob.sol, prob.clock, prob.vars, prob.params, prob.grid

  x, y = gridpoints(grid)

  # the Rossby wave initial condition
   ampl = 1e-2
  kwave = 3 * 2π/grid.Lx
  lwave = 2 * 2π/grid.Ly
      ω = -params.β * kwave / (kwave^2 + lwave^2 + 1 / deformation_radius^2)
    q₀  = @. ampl * cos(kwave * x) * cos(lwave * y)
    q₀h = rfft(q₀)

  SingleLayerQG.set_q!(prob, q₀)

  stepforward!(prob, nsteps)
  dealias!(sol, grid)
  SingleLayerQG.updatevars!(prob)

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

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

"""
    test_1layerqg_stochasticforcing_energy_budgets(dev; kwargs...)

Tests if the energy budget is closed for a SingleLayerQG problem with stochastic forcing.
"""
function test_1layerqg_stochasticforcing_energybudget(dev::Device=CPU(); n=256, dt=0.01, L=2π, ν=1e-7, nν=2, μ=1e-1, T=Float64)
  n, L  = 256, 2π
  ν, nν = 1e-7, 2
  μ = 1e-1
  dt, tf = 0.005, 0.1/μ
  nt = round(Int, tf/dt)

  grid = TwoDGrid(dev, n, L)
  x, y = gridpoints(grid)
  K  = @. sqrt(grid.Krsq)                      # a 2D array with the total wavenumber

  # Forcing with spectrum ~ exp(-(k-kf)²/(2 dkf²))
  kf, dkf = 12.0, 2.0
  ε = 0.1

  forcing_spectrum = zeros(dev, T, (grid.nkr, grid.nl))
  @. forcing_spectrum = exp(-(K - kf)^2 / (2 * dkf^2))
  @. forcing_spectrum = ifelse(grid.Krsq < 2^2,  0, forcing_spectrum)   # no power at low wavenumbers
  @. forcing_spectrum = ifelse(grid.Krsq > 20^2, 0, forcing_spectrum)   # no power at high wavenumbers
  ε0 = parsevalsum(forcing_spectrum .* grid.invKrsq / 2, grid) / (grid.Lx * grid.Ly)
  @. forcing_spectrum = ε / ε0 * forcing_spectrum

  Random.seed!(1234)

  function calcF!(Fh, sol, t, clock, vars, params, grid)
    eta = ArrayType(dev)(exp.(2π * im * rand(T, size(sol))) / sqrt(clock.dt))
    CUDA.@allowscalar eta[1, 1] = 0
    @. Fh = eta * sqrt(forcing_spectrum)

    return nothing
  end

  prob = SingleLayerQG.Problem(dev; nx=n, Lx=L, ν=ν, nν=nν, μ=μ, dt=dt,
   stepper="RK4", calcF=calcF!, stochastic=true)

  SingleLayerQG.set_q!(prob, 0*x)

  E = Diagnostic(SingleLayerQG.kinetic_energy,     prob, nsteps=nt)
  D = Diagnostic(SingleLayerQG.energy_dissipation, prob, nsteps=nt)
  R = Diagnostic(SingleLayerQG.energy_drag,        prob, nsteps=nt)
  W = Diagnostic(SingleLayerQG.energy_work,        prob, nsteps=nt)
  diags = [E, D, W, R]

  stepforward!(prob, diags, nt)

  SingleLayerQG.updatevars!(prob)

  dEdt_numerical = (E[2:E.i] - E[1:E.i-1]) / prob.clock.dt

  dEdt_computed = W[2:E.i] - D[1:E.i-1] - R[1:E.i-1]

  return isapprox(dEdt_numerical, dEdt_computed, rtol=1e-3)
end

"""
    test_1layerqg_deterministicforcing_energy_budgets(dev; kwargs...)

Tests if the energy budget is closed for a SingleLayerQG problem with deterministic forcing.
"""
function test_1layerqg_deterministicforcing_energybudget(dev::Device=CPU(); n=256, dt=0.01, L=2π, ν=1e-7, nν=2, μ=1e-1)
  n, L  = 256, 2π
  ν, nν = 1e-7, 2
  μ = 1e-1
  dt, tf = 0.005, 0.1/μ
  nt = round(Int, tf/dt)

  grid = TwoDGrid(dev, n, L)
  x, y = gridpoints(grid)
  k₀, l₀ = 2π / grid.Lx, 2π/grid.Ly

  # Forcing = 0.01cos(4x)cos(5y)cos(2t)
  f = @. 0.01 * cos(4k₀*x) * cos(5l₀*y)
  fh = rfft(f)

  function calcF!(Fh, sol, t, clock, vars, params, grid)
    @. Fh = fh * cos(2t)

    return nothing
  end

  prob = SingleLayerQG.Problem(dev; nx=n, Lx=L, ν=ν, nν=nν, μ=μ, dt=dt,
   stepper="RK4", calcF=calcF!, stochastic=false)

  SingleLayerQG.set_q!(prob, 0*x)

  E = Diagnostic(SingleLayerQG.kinetic_energy,     prob, nsteps=nt)
  D = Diagnostic(SingleLayerQG.energy_dissipation, prob, nsteps=nt)
  R = Diagnostic(SingleLayerQG.energy_drag,        prob, nsteps=nt)
  W = Diagnostic(SingleLayerQG.energy_work,        prob, nsteps=nt)
  diags = [E, D, W, R]

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

  SingleLayerQG.updatevars!(prob)

  dEdt_numerical = (E[3:E.i] - E[1:E.i-2]) / (2 * prob.clock.dt)
  dEdt_computed  = W[2:E.i-1] - D[2:E.i-1] - R[2:E.i-1]

  return isapprox(dEdt_numerical, dEdt_computed, rtol=1e-4)
end

"""
    test_1layerqg_stochasticforcing_enstrophy_budgets(dev; kwargs...)

Tests if the enstrophy budget is closed for a SingleLayerQG problem with stochastic forcing.
"""
function test_1layerqg_stochasticforcing_enstrophybudget(dev::Device=CPU(); n=256, dt=0.01, L=2π, ν=1e-7, nν=2, μ=1e-1, T=Float64)
  n, L  = 256, 2π
  ν, nν = 1e-7, 2
  μ = 1e-1
  dt, tf = 0.005, 0.1/μ
  nt = round(Int, tf/dt)

  grid = TwoDGrid(dev, n, L)
  x, y = gridpoints(grid)
  K  = @. sqrt(grid.Krsq)                      # a 2D array with the total wavenumber

  # Forcing with spectrum ~ exp(-(k-kf)²/(2 dkf²))
  kf, dkf = 12.0, 2.0
  εᶻ = 0.1

  forcing_spectrum = zeros(dev, T, (grid.nkr, grid.nl))
  @. forcing_spectrum = exp(-(K - kf)^2 / (2 * dkf^2))
  @. forcing_spectrum = ifelse(grid.Krsq < 2^2,  0, forcing_spectrum)   # no power at low wavenumbers
  @. forcing_spectrum = ifelse(grid.Krsq > 20^2, 0, forcing_spectrum)   # no power at high wavenumbers
  εᶻ0 = parsevalsum(forcing_spectrum / 2, grid) / (grid.Lx * grid.Ly)
  forcing_spectrum .= εᶻ / εᶻ0 * forcing_spectrum

  Random.seed!(1234)

  function calcF!(Fh, sol, t, clock, vars, params, grid)
    eta = ArrayType(dev)(exp.(2π * im * rand(T, size(sol))) / sqrt(clock.dt))
    CUDA.@allowscalar eta[1, 1] = 0
    @. Fh = eta * sqrt(forcing_spectrum)

    return nothing
  end

  prob = SingleLayerQG.Problem(dev; nx=n, Lx=L, ν=ν, nν=nν, μ=μ, dt=dt,
   stepper="RK4", calcF=calcF!, stochastic=true)

  SingleLayerQG.set_q!(prob, 0*x)

  Z = Diagnostic(SingleLayerQG.enstrophy,             prob, nsteps=nt)
  D = Diagnostic(SingleLayerQG.enstrophy_dissipation, prob, nsteps=nt)
  R = Diagnostic(SingleLayerQG.enstrophy_drag,        prob, nsteps=nt)
  W = Diagnostic(SingleLayerQG.enstrophy_work,        prob, nsteps=nt)
  diags = [Z, D, W, R]

  stepforward!(prob, diags, nt)

  SingleLayerQG.updatevars!(prob)

  dZdt_numerical = (Z[2:Z.i] - Z[1:Z.i-1]) / prob.clock.dt

  dZdt_computed = W[2:Z.i] - D[1:Z.i-1] - R[1:Z.i-1]

  return isapprox(dZdt_numerical, dZdt_computed, rtol=1e-3)
end

"""
    test_1layerqg_deterministicforcing_enstrophy_budgets(dev; kwargs...)

Tests if the enstrophy budget is closed for a SingleLayerQG problem with deterministic forcing.
"""
function test_1layerqg_deterministicforcing_enstrophybudget(dev::Device=CPU(); n=256, dt=0.01, L=2π, ν=1e-7, nν=2, μ=1e-1)
  n, L  = 256, 2π
  ν, nν = 1e-7, 2
  μ = 1e-1
  dt, tf = 0.005, 0.1/μ
  nt = round(Int, tf/dt)

  grid = TwoDGrid(dev, n, L)
  x, y = gridpoints(grid)
  k₀, l₀ = 2π/grid.Lx, 2π/grid.Ly

  # Forcing = 0.01cos(4x)cos(5y)cos(2t)
  f = @. 0.01 * cos(4k₀ * x) * cos(5l₀ * y)
  fh = rfft(f)

  function calcF!(Fh, sol, t, clock, vars, params, grid)
    @. Fh = fh * cos(2t)
    return nothing
  end

  prob = SingleLayerQG.Problem(dev; nx=n, Lx=L, ν=ν, nν=nν, μ=μ, dt=dt,
   stepper="RK4", calcF=calcF!, stochastic=false)

  SingleLayerQG.set_q!(prob, 0*x)

  Z = Diagnostic(SingleLayerQG.enstrophy,             prob, nsteps=nt)
  D = Diagnostic(SingleLayerQG.enstrophy_dissipation, prob, nsteps=nt)
  R = Diagnostic(SingleLayerQG.enstrophy_drag,        prob, nsteps=nt)
  W = Diagnostic(SingleLayerQG.enstrophy_work,        prob, nsteps=nt)
  diags = [Z, D, W, R]

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

  SingleLayerQG.updatevars!(prob)

  dZdt_numerical = (Z[3:Z.i] - Z[1:Z.i-2]) / (2 * prob.clock.dt)
  dZdt_computed  = W[2:Z.i-1] - D[2:Z.i-1] - R[2:Z.i-1]

  return isapprox(dZdt_numerical, dZdt_computed, rtol=1e-4)
end

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

Tests the advection term in the SingleLayerQG 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) - ν∇²ζ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_1layerqg_advection(dt, stepper, dev::Device=CPU(); n=128, L=2π, ν=1e-2, nν=1, μ=0.0)
  n, L  = 128, 2π
  ν, nν = 1e-2, 1
   μ = 0.0
  tf = 1.0
  nt = round(Int, tf/dt)

  grid = TwoDGrid(dev, n, L)
  x, y = gridpoints(grid)

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

  Ff = @. -(
    ν*( 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)

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

  prob = SingleLayerQG.Problem(dev; nx=n, Lx=L, ν=ν, nν=nν, μ=μ, dt=dt, stepper=stepper, calcF=calcF!)

  SingleLayerQG.set_q!(prob, qf)

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

  SingleLayerQG.updatevars!(prob)

  return isapprox(prob.vars.q, qf, rtol=rtol_singlelayerqg)
end

"""
    test_1layerqg_energyenstrophy_BarotropicQG(dev)

Tests the energy and enstrophy function for a SingleLayerQG problem.
"""
function test_1layerqg_energyenstrophy_BarotropicQG(dev::Device=CPU())
  nx, Lx  = 64, 2π
  ny, Ly  = 64, 3π
  grid = TwoDGrid(dev, nx, Lx, ny, Ly)
  k₀, l₀ = 2π/Lx, 2π/Ly # fundamental wavenumbers
  x, y = gridpoints(grid)

  energy_calc = 29/9
  enstrophy_calc = 10885/648

  η  = @. cos(10k₀ * x) * cos(10l₀ * y)
  ψ₀ = @. sin(2k₀ * x) * cos(2l₀ * y) + 2sin(k₀ * x) * cos(3l₀ * y)
  q₀ = @. - ((2k₀)^2 + (2l₀)^2) * sin(2k₀ * x) * cos(2l₀ * y) - (k₀^2 + (3l₀)^2) * 2sin(k₀ * x) * cos(3l₀*y)

  prob = SingleLayerQG.Problem(dev; nx=nx, Lx=Lx, ny=ny, Ly=Ly, eta=η, stepper="ForwardEuler")
  sol, clock, vars, params, grid = prob.sol, prob.clock, prob.vars, prob.params, prob.grid
  SingleLayerQG.set_q!(prob, q₀)
  SingleLayerQG.updatevars!(prob)

  energyq₀ = SingleLayerQG.energy(prob)
  enstrophyq₀ = SingleLayerQG.enstrophy(prob)

  return isapprox(energyq₀, energy_calc, rtol=rtol_singlelayerqg) && isapprox(enstrophyq₀, enstrophy_calc, rtol=rtol_singlelayerqg) && SingleLayerQG.potential_energy(prob)==0 &&
  SingleLayerQG.addforcing!(prob.timestepper.N, sol, clock.t, clock, vars, params, grid) == nothing
end

"""
    test_1layerqg_energies_EquivalentBarotropicQG(dev)

Tests the kinetic and potential energy for an equivalent barotropic SingleLayerQG problem.
"""
function test_1layerqg_energies_EquivalentBarotropicQG(dev; deformation_radius=1.23)
  nx, Lx  = 64, 2π
  ny, Ly  = 64, 3π
  grid = TwoDGrid(dev, nx, Lx, ny, Ly)
  k₀, l₀ = 2π/Lx, 2π/Ly # fundamental wavenumbers
  x, y = gridpoints(grid)

  kinetic_energy_calc = 29/9
  potential_energy_calc = 5/(8*deformation_radius^2)
  energy_calc = kinetic_energy_calc + potential_energy_calc

  η  = @. cos(10k₀ * x) * cos(10l₀ * y)
  ψ₀ = @. sin(2k₀ * x) * cos(2l₀ * y) + 2sin(k₀ * x) * cos(3l₀ * y)
  q₀ = @. - ((2k₀)^2 + (2l₀)^2) * sin(2k₀ * x) * cos(2l₀ * y) - (k₀^2 + (3l₀)^2) * 2sin(k₀ * x) * cos(3l₀*y) - 1/deformation_radius^2 * ψ₀

  prob = SingleLayerQG.Problem(dev; nx=nx, Lx=Lx, ny=ny, Ly=Ly, eta=η, deformation_radius=deformation_radius, stepper="ForwardEuler")
  sol, clock, vars, params, grid = prob.sol, prob.clock, prob.vars, prob.params, prob.grid
  SingleLayerQG.set_q!(prob, q₀)
  SingleLayerQG.updatevars!(prob)

  kinetic_energyq₀ = SingleLayerQG.kinetic_energy(prob)
  potential_energyq₀ = SingleLayerQG.potential_energy(prob)
  energyq₀ = SingleLayerQG.energy(prob)

  return isapprox(kinetic_energyq₀, kinetic_energy_calc, rtol=rtol_singlelayerqg) && isapprox(potential_energyq₀, potential_energy_calc, rtol=rtol_singlelayerqg) && isapprox(energyq₀, energy_calc, rtol=rtol_singlelayerqg) &&
  SingleLayerQG.addforcing!(prob.timestepper.N, sol, clock.t, clock, vars, params, grid) == nothing
end

"""
    test_1layerqg_problemtype(dev, T; deformation_radius=Inf)

Tests the SingleLayerQG problem constructor for different DataType `T`.
"""
function test_1layerqg_problemtype(dev, T; deformation_radius=Inf)
  prob = SingleLayerQG.Problem(dev; T=T, deformation_radius=deformation_radius)

  A = ArrayType(dev)

  return (typeof(prob.sol)<:A{Complex{T}, 2} && typeof(prob.grid.Lx)==T && eltype(prob.grid.x)==T && typeof(prob.vars.u)<:A{T, 2})
end

function test_streamfunctionfrompv(dev; deformation_radius=1.23)
  prob_barotropicQG = SingleLayerQG.Problem(dev; nx=64, deformation_radius=Inf)
  prob_equivalentbarotropicQG = SingleLayerQG.Problem(dev; nx=64, deformation_radius=deformation_radius)

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

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

  q_barotropic = @. -((2k₀)^2 + (3l₀)^2) * sin(2k₀ * x) * cos(3l₀ * y) - 0*(3k₀)^2 * sin(3k₀ * x)
  q_equivalentbarotropic = @. -((2k₀)^2 + (3l₀)^2 + 1/deformation_radius^2) * sin(2k₀ * x) * cos(3l₀ * y) - 0*((3k₀)^2+ 1/deformation_radius^2) * sin(3k₀ * x)

  SingleLayerQG.set_q!(prob_barotropicQG, q_barotropic)
  SingleLayerQG.set_q!(prob_equivalentbarotropicQG, q_equivalentbarotropic)

  SingleLayerQG.streamfunctionfrompv!(prob_barotropicQG.vars.ψh, prob_barotropicQG.vars.qh, prob_barotropicQG.params, prob_barotropicQG.grid)

  SingleLayerQG.streamfunctionfrompv!(prob_equivalentbarotropicQG.vars.ψh, prob_equivalentbarotropicQG.vars.qh, prob_equivalentbarotropicQG.params, prob_equivalentbarotropicQG.grid)

  return (prob_barotropicQG.vars.ψ ≈ ψ && prob_equivalentbarotropicQG.vars.ψ ≈ ψ)
end

function test_1layerqg_energy_dissipation(dev; deformation_radius=2.23)
  prob = SingleLayerQG.Problem(dev; nx=16, deformation_radius=deformation_radius)
  SingleLayerQG.energy_dissipation(prob)
  return nothing
end

function test_1layerqg_enstrophy_dissipation(dev; deformation_radius=2.23)
  prob = SingleLayerQG.Problem(dev; nx=16, deformation_radius=deformation_radius)
  SingleLayerQG.enstrophy_dissipation(prob)
  return nothing
end

function test_1layerqg_energy_work(dev; deformation_radius=2.23)
  prob = SingleLayerQG.Problem(dev; nx=16, deformation_radius=deformation_radius)
  SingleLayerQG.energy_work(prob)
  return nothing
end

function test_1layerqg_enstrophy_work(dev; deformation_radius=2.23)
  prob = SingleLayerQG.Problem(dev; nx=16, deformation_radius=deformation_radius)
  SingleLayerQG.enstrophy_work(prob)
  return nothing
end

function test_1layerqg_energy_drag(dev; deformation_radius=2.23)
  prob = SingleLayerQG.Problem(dev; nx=16, deformation_radius=deformation_radius)
  SingleLayerQG.energy_drag(prob)
  return nothing
end

function test_1layerqg_enstrophy_drag(dev; deformation_radius=2.23)
  prob = SingleLayerQG.Problem(dev; nx=16, deformation_radius=deformation_radius)
  SingleLayerQG.enstrophy_drag(prob)
  return nothing
end


module RadauBig

using ForwardDiff
using StaticArrays
using LinearAlgebra
using DelimitedFiles
using Polynomials: Poly, polyval, polyder, coeffs
using PolynomialRoots: roots
using GenericLinearAlgebra: eigvals
using GenericSVD: svd

include("big.jl")
include("generate_butcher_table.jl")
include("save_radau_tables.jl")

export
    # big.jl
    big_eigen,

    # generate_butcher_table.jl
    find_real_eigenvalue,
    radau_butcher_table_plus,

    # save_radau_tables.jl
    save_radau_table_to_file

end


using Test
using ModelingToolkit, OrdinaryDiffEq

@parameters t
@connector function Pin(;name)
    sts = @variables v(t)=1.0 i(t)=1.0
    ODESystem(Equation[], t, sts, []; name=name)
end

function ModelingToolkit.connect(::Type{Pin}, ps...)
    eqs = [
           0 ~ sum(p->p.i, ps) # KCL
          ]
    # KVL
    for i in 1:length(ps)-1
        push!(eqs, ps[i].v ~ ps[i+1].v)
    end

    return eqs
end

function Ground(;name)
    @named g = Pin()
    eqs = [g.v ~ 0]
    compose(ODESystem(eqs, t, [], []; name=name), g)
end

function OnePort(;name)
    @named p = Pin()
    @named n = Pin()
    sts = @variables v(t)=1.0 i(t)=1.0
    eqs = [
           v ~ p.v - n.v
           0 ~ p.i + n.i
           i ~ p.i
          ]
    compose(ODESystem(eqs, t, sts, []; name=name), p, n)
end

function Resistor(;name, R = 1.0)
    @named oneport = OnePort()
    @unpack v, i = oneport
    ps = @parameters R=R
    eqs = [
           v ~ i * R
          ]
    extend(ODESystem(eqs, t, [], ps; name=name), oneport)
end

function Capacitor(;name, C = 1.0)
    @named oneport = OnePort()
    @unpack v, i = oneport
    ps = @parameters C=C
    D = Differential(t)
    eqs = [
           D(v) ~ i / C
          ]
    extend(ODESystem(eqs, t, [], ps; name=name), oneport)
end

function ConstantVoltage(;name, V = 1.0)
    @named oneport = OnePort()
    @unpack v = oneport
    ps = @parameters V=V
    eqs = [
           V ~ v
          ]
    extend(ODESystem(eqs, t, [], ps; name=name), oneport)
end

function Inductor(; name, L = 1.0)
    @named oneport = OnePort()
    @unpack v, i = oneport
    ps = @parameters L=L
    D = Differential(t)
    eqs = [
           D(i) ~ v / L
          ]
    extend(ODESystem(eqs, t, [], ps; name=name), oneport)
end


module FEM

using LinearAlgebra, TensorRefinement.Auxiliary, ..TensorTrain

export extdn, diffextdn, diffdn, intdn, bpxdn, diffbpxdn

function extdn(::Type{T}, L::Int, ℓ::Int, d::Int; major::String="last") where {T<:FloatRC}
	if L < 1
		throw(ArgumentError("L should be at least one"))
	end
	if ℓ ∉ 0:L
		throw(ArgumentError("ℓ should be in 0:L"))
	end
	if d ≤ 0
		throw(ArgumentError("d should be positive"))
	end
	if major ∉ ("first", "last")
		throw(ArgumentError("major should be either \"last\" (default) or \"first\""))
	end
	I = [1 0; 0 1]; J = [0 1; 0 0]; O = [0 0; 0 0]
	U = Matrix{T}([I J'; O J]); U = factor(U, 2, 2)
	W = Matrix{T}([
		1 0
		2 1
		1 2
		0 1
		] / 2 / sqrt(2)); W = factor(W, 2, 1)
	U = factorkp((U,d))
	W = factorkp((W,d))
	if ℓ == 0
		P = dec(W, len=L-ℓ)
	else
		P = dec(U, len=ℓ)
		(ℓ < L) && decappend!(P, dec(W, len=L-ℓ))
	end
	decrankselect!(P, 1:1, 1:1)
	(major == "last") && decreverse!(P)
	return P
end

extdn(L::Int, ℓ::Int, d::Int; major::String="last") = extdn(Float64, L, ℓ, d; major=major)

function diffdn(::Type{T}, L::Int, d::Int, k::Int; major::String="last") where {T<:FloatRC}
	if L < 1
		throw(ArgumentError("L should be at least one"))
	end
	if d < 0
		throw(ArgumentError("d should be positive"))
	end
	if k ∉ 0:d
		throw(ArgumentError("k should be in 0:d"))
	end
	if major ∉ ("first", "last")
		throw(ArgumentError("major should be either \"last\" (default) or \"first\""))
	end
	I = [1 0; 0 1]; J = [0 1; 0 0]; O = [0 0; 0 0]
	S = Matrix{T}([1 1; 1 -1]); S = factor(S, 1, 1)
	W₀ = Matrix{T}([I+J+J' I-J-J'; I-J+J' I+J-J'] / sqrt(2)); W₀ = factor(W₀, 2, 2)
	W₁ = W₀ * 2
	Z₀ = Matrix{T}([1; 0; 0; 1][:,:] / 2); Z₀ = factor(Z₀, 2, 1)
	Z₁ = Matrix{T}([0; 1][:,:]); Z₁ = factor(Z₁, 1, 1)
	S = factorkp((S,d))
	W = factorkp([ (k == i) ? W₁ : W₀ for i ∈ 1:d ]...)
	Z = factorkp([ (k == i) ? Z₁ : Z₀ for i ∈ 1:d ]...)
	M = dec(S)
	decappend!(M, dec(W, len=L))
	decpush!(M, Z)
	decrankselect!(M, 1:1, :)
	decskp!(M, 1, path="forward")
	(major == "last") && decreverse!(M)
	return M
end

diffdn(L::Int, d::Int, k::Int; major::String="last") = diffdn(Float64, L, d, k; major=major)

function intdn(::Type{T}, L::Int, d::Int, kk::Int, k::Int; major::String="last") where {T<:FloatRC}
	if L < 1
		throw(ArgumentError("L should be at least one"))
	end
	if d < 1
		throw(ArgumentError("d should be positive"))
	end
	if (kk,k) ⊈ 0:d
		throw(ArgumentError("kk and k should be in 0:d"))
	end
	if major ∉ ("first", "last")
		throw(ArgumentError("major should be either \"last\" (default) or \"first\""))
	end
	I = [1 0; 0 1]
	U = Matrix{T}(I/2); U = factor(U, 2, 2)
	Z = Matrix{T}([2 0; 0 2/3] / 2); Z = factor(Z, 2, 2)

	U = factorkp((U,d))
	Z = [ Z[:,1:((i == kk) ? 1 : 2),1:((i == k) ? 1 : 2),:] for i ∈ 1:d ]
	Z = factorkp(Z...)

	Λ = dec(U, len=L)
	decpush!(Λ, Z)
	(major == "last") && decreverse!(Λ)
	return Λ
end

intdn(L::Int, d::Int, kk::Int, k::Int; major::String="last") = intdn(Float64, L, d, kk, k; major=major)

function bpxdn(::Type{T}, L::Int, d::Int; major::String="last", κ::T=one(T), balance::Bool=true) where {T<:FloatRC}
	if L < 1
		throw(ArgumentError("L should be positive"))
	end
	if d < 1
		throw(ArgumentError("d should be positive"))
	end
	if major ∉ ("first", "last")
		throw(ArgumentError("major should be either \"last\" (default) or \"first\""))
	end
	I = [1 0; 0 1]; J = [0 1; 0 0]; O = [0 0; 0 0]
	U = Matrix{T}([I J'; O J]); U = factor(U, 2, 2)
	W = Matrix{T}([
		1 0
		2 1
		1 2
		0 1
		] / 2 / sqrt(2)); W = factor(W, 2, 1)
	U = factorkp((factormp(U, 2, U, 2),d))
	W = factorkp((factormp(W, 2, W, 2),d))
	μ = 2.0^(κ/2)
	if L == 1
		return dec(U/μ^2 + W)
	end
	if balance
		C = dec(factorhcat(U/μ, U/μ + W*μ))
		for ℓ ∈ 2:L-1
			decpush!(C, factorutcat(U/μ, U/μ, W*μ))
		end
		decpush!(C, factorvcat(U/μ^(L+1), W/μ^(L-1)))
	else
		C = dec(factorhcat(U, U/μ^2 + W))
		for ℓ ∈ 2:L-1
			decpush!(C, factorutcat(U, U/μ^(2*ℓ), W))
		end
		decpush!(C, factorvcat(U/μ^(2*L), W))
	end
	if major == "last"
		decreverse!(C)
	end
	return C
end

bpxdn(L::Int, d::Int; major::String="last", κ::Float64=Float64(1), balance::Bool=true) = bpxdn(Float64, L, d; major=major, κ=κ, balance=balance)

function diffextdn(::Type{T}, L::Int, ℓ::Int, d::Int, k::Int; major::String="last") where {T<:FloatRC}
	if L < 1
		throw(ArgumentError("L should be at least one"))
	end
	if ℓ ∉ 0:L
		throw(ArgumentError("ℓ should be in 0:L"))
	end
	if d < 1
		throw(ArgumentError("d should be positive"))
	end
	if k ∉ 0:d
		throw(ArgumentError("k should be in 0:d"))
	end
	if major ∉ ("first", "last")
		throw(ArgumentError("major should be either \"last\" (default) or \"first\""))
	end
	I = [1 0; 0 1]; J = [0 1; 0 0]; O = [0 0; 0 0]

	A  = Matrix{T}([1 0]); A = factor(A, 1, 1)
	S  = Matrix{T}([1 1; 1 -1]); S = factor(S, 1, 1)
	X  = factorcontract(A, S)
	V₀ = Matrix{T}([I+J+J' I-J-J'; I-J+J' I+J-J'] / sqrt(2)); V₀ = factor(V₀, 2, 2)
	V₁ = V₀ * 2
	R₀ = Matrix{T}([1 0; 0 1]); R₀ = factor(R₀, 1, 1)
	R₁ = Matrix{T}([0; 1][:,:]); R₁ = factor(R₁, 1, 1)
	Y₀ = Matrix{T}([
		 2 0
		 2 0
		-1 1
		 1 1
		 ] / 2); Y₀ = factor(Y₀, 2, 1)
	Y₁ = Y₀[2:2,:,:,2:2] * 2
	Z₀ = Matrix{T}([1,0,0,1][:,:] / 2); Z₀ = factor(Z₀, 2, 1)
	Z₁ = Matrix{T}([1][:,:]); Z₁ = factor(Z₁, 1, 1)

	X = factorkp((X, d))
	V = factorkp([ (k == i) ? V₁ : V₀ for i ∈ 1:d ]...)
	R = factorkp([ (k == i) ? R₁ : R₀ for i ∈ 1:d ]...)
	Y = factorkp([ (k == i) ? Y₁ : Y₀ for i ∈ 1:d ]...)
	Z = factorkp([ (k == i) ? Z₁ : Z₀ for i ∈ 1:d ]...)

	XR = factorcontract(X, R)
	VR = factorcontract(V, R)
	if ℓ == 0
		Q = dec(XR)
	else
		Q = dec(X)
		if ℓ > 1
			decappend!(Q, dec(V, len=ℓ-1))
		end
		decpush!(Q, VR)
	end
	if ℓ < L
		decappend!(Q, dec(Y, len=L-ℓ))
	end
	decpush!(Q, Z)
	decskp!(Q, 1, path="forward")
	(major == "last") && decreverse!(Q)
	return Q
end

diffextdn(L::Int, ℓ::Int, d::Int, k::Int; major::String="last") = diffextdn(Float64, L, ℓ, d, k; major=major)

function diffbpxdn(::Type{T}, L::Int, d::Int, k::Int; major::String="last") where {T<:FloatRC}
	if L < 1
		throw(ArgumentError("L should be at least one"))
	end
	if d < 1
		throw(ArgumentError("d should be positive"))
	end
	if k ∉ 0:d
		throw(ArgumentError("k should be in 0:d"))
	end
	if major ∉ ("first", "last")
		throw(ArgumentError("major should be either \"last\" (default) or \"first\""))
	end
	I = [1 0; 0 1]; J = [0 1; 0 0]; O = [0 0; 0 0]

	A  = Matrix{T}([1 0]); A = factor(A, 1, 1)
	S  = Matrix{T}([1 1; 1 -1]); S = factor(S, 1, 1)
	X  = factorcontract(A, S)
	V₀ = Matrix{T}([I+J+J' I-J-J'; I-J+J' I+J-J'] / sqrt(2)); V₀ = factor(V₀, 2, 2)
	V₁ = V₀ * 2
	R₀ = Matrix{T}([1 0; 0 1]); R₀ = factor(R₀, 1, 1)
	R₁ = Matrix{T}([0; 1][:,:]); R₁ = factor(R₁, 1, 1)
	Y₀ = Matrix{T}([
		 2 0
		 2 0
		-1 1
		 1 1
		 ] / 2); Y₀ = factor(Y₀, 2, 1)
	Y₁ = Y₀[2:2,:,:,2:2] * 2
	Z₀ = Matrix{T}([1,0,0,1][:,:] / 2); Z₀ = factor(Z₀, 2, 1)
	Z₁ = Matrix{T}([1][:,:]); Z₁ = factor(Z₁, 1, 1)

	U = Matrix{T}([I J'; O J]); U = factor(U, 2, 2)
	R = Matrix{T}([1 0; 0 1]); R = factor(R, 1, 1)
	W = Matrix{T}([
		1 0
		2 1
		1 2
		0 1
		] / 2 / sqrt(2)); W = factor(W, 2, 1)
	Z = Matrix{T}([1; 0][:,:]); Z = factor(Z, 1, 1)

	X  = factormp(X, 2, A, 2)
	V₀ = factormp(V₀, 2, U, 2)
	V₁ = factormp(V₁, 2, U, 2)
	R₀ = factormp(R₀, 2, R, 2)
	R₁ = factormp(R₁, 2, R, 2)
	Y₀ = factormp(Y₀, 2, W, 2)
	Y₁ = factormp(Y₁, 2, W, 2)
	Z₀ = factormp(Z₀, 2, Z, 2)
	Z₁ = factormp(Z₁, 2, Z, 2)

	X  = factorkp((X, d))
	V  = factorkp([ (k == i) ? V₁ : V₀ for i ∈ 1:d ]...)
	R  = factorkp([ (k == i) ? R₁ : R₀ for i ∈ 1:d ]...)
	Y  = factorkp([ (k == i) ? Y₁ : Y₀ for i ∈ 1:d ]...)
	Z  = factorkp([ (k == i) ? Z₁ : Z₀ for i ∈ 1:d ]...)

	XR = factorcontract(X, R)
	VR = factorcontract(V, R)
	μ = 2*one(T)
	Q = dec(factorhcat(X, XR))
	for ℓ ∈ 1:L-1
		decpush!(Q, factorutcat(V/μ, VR/μ, Y))
	end
	decpush!(Q, factorvcat(VR/μ, Y))
	decpush!(Q, Z)
	decskp!(Q, 1, path="forward")
	(major == "last") && decreverse!(Q)
	return Q
end

diffbpxdn(L::Int, d::Int, k::Int; major::String="last") = diffbpxdn(Float64, L, d, k; major=major)

end


using TensorOperators
using Base.Test
using ArrayViews

input_framesize = 5
output_framesize = 4
kW = 2
dW = 1

m = TemporalConvolution(cpu, Float64, input_framesize, output_framesize, kW, dW)
init!(m)

x = rand(5, 10)

y = forward!(m, x)

# Compute the output

output_length = div(size(x, 2) - kW, dW) + 1
y2 = zeros(output_framesize, output_length)
for i in 1:output_length
    xi = flatten_view(x[:, i:(i+kW-1)])
    y2[:, i] = m.bias.values + m.weight.values * xi
end

@test_approx_eq y y2