# “One size fits all” analytic solutions to the Grad–Shafranov equation
# Phys. Plasmas 17, 032502 (2010); https://doi.org/10.1063/1.3328818
#
# Ideal MHD Freidberg Chapter 6.6.1

@fastmath function _solovev_polynomials(x, y, ::Val{N} = Val(12); logx=log(x)) where N

    x2 = x^2
    x4 = x2^2
    x6 = x2^3
    y2 = y^2
    y4 = y2^2
    y6 = y2^3

    # Eq. 8 & Eq. 27
    ψ₁  = 1
    ψ₂  = x2
    ψ₃  = y2 - x2*logx
    ψ₄  = x4 - 4*x2*y2
    ψ₅  = 2*y4 - 9*y2*x2 + 3*x4*logx - 12*x2*y2*logx
    ψ₆  = x6 - 12*x4*y2 + 8*x2*y4
    ψ₇  = 8*y6 - 140*y4*x2 + 75*y2*x4 - 15*x6*logx + 180*x4*y2*logx - 120*x2*y4*logx

    N == 7 && return @SVector [ψ₁, ψ₂, ψ₃, ψ₄, ψ₅, ψ₆, ψ₇]

    x3 = x^3
    y3 = y^3
    y5 = y^5

    ψ₈  = y
    ψ₉  = y*x2
    ψ₁₀ = y3 - 3*y*x2*logx
    ψ₁₁ = 3*y*x4 - 4*y3*x2
    ψ₁₂ = 8*y5 - 45*y*x4 - 80*y3*x2*logx + 60*y*x4*logx

    ψ_polys = @SVector [ψ₁, ψ₂, ψ₃, ψ₄, ψ₅, ψ₆, ψ₇, ψ₈, ψ₉, ψ₁₀, ψ₁₁, ψ₁₂]

    return ψ_polys
end

@fastmath function _solovev_polynomials_Dx(x, y, ::Val{N} = Val(12); logx=log(x)) where N

    x2 = x^2
    x3 = x^3
    x4 = x2^2
    x5 = x^5
    x6 = x2^3
    y2 = y^2
    y4 = y2^2

    # Derivatives from Mathematica
    ψ₁  = zero(x)
    ψ₂  = 2*x
    ψ₃  = -(2*x*logx+ x)
    ψ₄  = 4*x3 - 8*x*y2
    ψ₅  = 3*x3 - 30*x*y2 + 12*x3*logx - 24*x*y2*logx
    ψ₆  = 6*x5 - 48*x3*y2 + 16*x*y4
    ψ₇  = -5*x*(3*x4 - 96*x2*y2 + 80*y4 + 6*(3*x4 - 24*x2*y2 + 8*y4)*logx)

    N == 7 && return @SVector [ψ₁, ψ₂, ψ₃, ψ₄, ψ₅, ψ₆, ψ₇]

    y3 = y^3

    ψ₈  = zero(x)
    ψ₉  = 2*y*x
    ψ₁₀ = -3*x*y*(1 + 2*logx)
    ψ₁₁ = 12*x3*y - 8*x*y3
    ψ₁₂ = 40*x*y*(-3*x2 - 2*y2 + (6*x2 - 4*y2)*logx)

    ψ_polys = @SVector [ψ₁, ψ₂, ψ₃, ψ₄, ψ₅, ψ₆, ψ₇, ψ₈, ψ₉, ψ₁₀, ψ₁₁, ψ₁₂]

    return ψ_polys
end

@fastmath function _solovev_polynomials_Dxx(x, y, ::Val{N} = Val(12); logx=log(x)) where N

    x2 = x^2
    x3 = x^3
    x4 = x2^2
    x6 = x2^3
    y2 = y^2
    y4 = y2^2

    # Derivatives from Mathematica
    ψ₁  = zero(x)
    ψ₂  = 2*one(x)
    ψ₃  = -2*logx - 3
    ψ₄  = 12*x2 - 8*y2
    ψ₅  = 3*(7*x2 - 18*y2 + 4*(3*x2 - 2*y2)*logx)
    ψ₆  = 2*(15*x4 - 72*x2*y2 + 8*y4)
    ψ₇  = -5*(33*x4 - 432*x2*y2 + 128*y4 + 6*(15*x4 - 72*x2*y2 + 8*y4)*logx)

    N == 7 && return @SVector [ψ₁, ψ₂, ψ₃, ψ₄, ψ₅, ψ₆, ψ₇]

    y3 = y^3

    ψ₈  = zero(x)
    ψ₉  = 2*y
    ψ₁₀ = -3*y*(3 + 2*logx)
    ψ₁₁ = 36*x2*y - 8*y3
    ψ₁₂ = -40*y*(3*x2 + 6*y2 - 18*x2*logx + 4*y2*logx)

    ψ_polys = @SVector [ψ₁, ψ₂, ψ₃, ψ₄, ψ₅, ψ₆, ψ₇, ψ₈, ψ₉, ψ₁₀, ψ₁₁, ψ₁₂]

    return ψ_polys
end

@fastmath function _solovev_polynomials_Dy(x, y, ::Val{N} = Val(12); logx=log(x)) where N

    x2 = x^2
    x3 = x^3
    x4 = x2^2
    x6 = x2^3
    y2 = y^2
    y3 = y^3
    y4 = y2^2

    # Derivatives from Mathematica
    ψ₁  = zero(y)
    ψ₂  = zero(y)
    ψ₃  = 2*y
    ψ₄  = -8*x2*y
    ψ₅  = 2*y*(-9*x2 + 4*y2 - 12*x2*logx)
    ψ₆  = -24*x4*y + 32*x2*y3
    ψ₇  = 2*y*(75*x4 - 280*x2*y2 + 24*y4 + 60*(3*x4 - 4*x2*y2)*logx)

    N == 7 && return @SVector [ψ₁, ψ₂, ψ₃, ψ₄, ψ₅, ψ₆, ψ₇]

    ψ₈  = one(y)
    ψ₉  = x2
    ψ₁₀ = 3*(y2 - x2*logx)
    ψ₁₁ = 3*(x4 - 4*x2*y2)
    ψ₁₂ = 5*(-9*x4 + 8*y4 + 12*(x4 - 4*x2*y2)*logx)

    ψ_polys = @SVector [ψ₁, ψ₂, ψ₃, ψ₄, ψ₅, ψ₆, ψ₇, ψ₈, ψ₉, ψ₁₀, ψ₁₁, ψ₁₂]

    return ψ_polys
end

@fastmath function _solovev_polynomials_Dyy(x, y, ::Val{N} = Val(12); logx=log(x)) where N

    x2 = x^2
    x3 = x^3
    x4 = x2^2
    x6 = x2^3
    y2 = y^2
    y3 = y^3
    y4 = y2^2

    # Derivatives from Mathematica
    ψ₁  = zero(y)
    ψ₂  = zero(y)
    ψ₃  = 2*one(y)
    ψ₄  = -8*x2
    ψ₅  = -6*(3*x2 - 4*y2 + 4*x2*logx)
    ψ₆  = -24*x2*(x2 - 4*y2)
    ψ₇  = 30*(5*x4 - 56*x2*y2 + 8*y4 + 12*(x4 - 4*x2*y2)*logx)

    N == 7 && return @SVector [ψ₁, ψ₂, ψ₃, ψ₄, ψ₅, ψ₆, ψ₇]

    ψ₈  = zero(y)
    ψ₉  = zero(y)
    ψ₁₀ = 6*y
    ψ₁₁ = -24*x2*y
    ψ₁₂ = 160*y*(y2 - 3*x2*logx)

    ψ_polys = @SVector [ψ₁, ψ₂, ψ₃, ψ₄, ψ₅, ψ₆, ψ₇, ψ₈, ψ₉, ψ₁₀, ψ₁₁, ψ₁₂]

    return ψ_polys
end

@fastmath function _solovev_psi_p(x,y,A; logx=log(x))
    # Derivatives from Mathematica
    x2 = x^2
    x4 = x2*x2
    return x4/8 + A*(x2*logx/2 - x4/8)
end

@fastmath function _solovev_psi_p_Dx(x,y,A; logx=log(x))
    # Derivatives from Mathematica
    x2 = x^2
    return (x*(A + x2 - A*x2 + 2*A*logx))/2
end

@fastmath function _solovev_psi_p_Dxx(x,y,A; logx=log(x))
    # Derivatives from Mathematica
    x2 = x^2
    return (3*(A + x2 - A*x2))/2 + A*logx
end

function _solovev_psi_p_Dy(x,y,A)
    return zero(y)
end

function _solovev_psi_p_Dyy(x,y,A)
    return zero(y)
end

function _solovev_psi(x,y,A,coeffs::SVector{N};logx=log(x)) where N

    ψ_p = _solovev_psi_p(x,y,A; logx=logx)
    ψ_polys = _solovev_polynomials(x,y,Val(N),logx=logx)

    ψ_h = zero(x)
    @inbounds for i=1:N
        ψ_h += coeffs[i]*ψ_polys[i]
    end

    return ψ_p + ψ_h
end

function _solovev_psi_Dx(x,y,A,coeffs::SVector{N};logx=log(x)) where N

    ψ_p = _solovev_psi_p_Dx(x,y,A,logx=logx)
    ψ_polys = _solovev_polynomials_Dx(x,y,Val(N),logx=logx)

    ψ_h = zero(x)
    @inbounds for i=1:N
        ψ_h += coeffs[i]*ψ_polys[i]
    end

    return ψ_p + ψ_h
end

function _solovev_psi_Dxx(x,y,A,coeffs::SVector{N};logx=log(x)) where N

    ψ_p = _solovev_psi_p_Dxx(x,y,A,logx=logx)
    ψ_polys = _solovev_polynomials_Dxx(x,y,Val(N),logx=logx)

    ψ_h = zero(x)
    @inbounds for i=1:N
        ψ_h += coeffs[i]*ψ_polys[i]
    end

    return ψ_p + ψ_h
end

function _solovev_psi_Dy(x,y,A,coeffs::SVector{N};logx=log(x)) where N

    ψ_p = _solovev_psi_p_Dy(x,y,A)
    ψ_polys = _solovev_polynomials_Dy(x,y,Val(N),logx=logx)

    ψ_h = zero(y)
    @inbounds for i=1:N
        ψ_h += coeffs[i]*ψ_polys[i]
    end

    return ψ_p + ψ_h
end

function _solovev_psi_Dyy(x,y,A,coeffs::SVector{N};logx=log(x)) where N

    ψ_p = _solovev_psi_p_Dyy(x,y,A)
    ψ_polys = _solovev_polynomials_Dyy(x,y,Val(N),logx=logx)

    ψ_h = zero(x)
    @inbounds for i=1:N
        ψ_h += coeffs[i]*ψ_polys[i]
    end

    return ψ_p + ψ_h
end

"""
SolovevEquilibrium Structure

Defines the equilibrium that satisfy Δ⋆ψ(x,y) = α + (1 - α)x^2\\
where (r,z) = (R0 x,R0 y). F(dF/dψ) = -A, μ₀dp/dψ = -C, α = A/(A + C R0²), ψ₀ = R0²(A + C R0²)

Fields:\\
`cocos` - COCOS\\
`B0` - Toroidal magnetic field magnitude on axis [T]\\
`R0` - Major Radius [m]\\
`epsilon` - inverse aspect ratio a/R0\\
`delta` - triangularity\\
`kappa` - elongation/ellipticity\\
`alpha` - constant relating beta regime (α)\\
`qstar` - Kink safety factor\\
`psi0` - Poloidal flux normalization (ψ₀)\\
`beta_p` - Poloidal beta\\
`beta_t` - Toroidal beta\\
`c` - Coefficients for Solov'ev polynomials\\
`diverted` - If true then equilibrium has one or more x-points\\
`symmetric` - If true then equilibrium is up-down symmetric\\
`sigma` - Sign of dot(B,J)
"""
struct SolovevEquilibrium{T,N} <: AbstractEquilibrium
    cocos::COCOS
    B0::T
    R0::T
    epsilon::T
    delta::T
    kappa::T
    alpha::T
    qstar::T
    psi0::T
    beta_p::T
    beta_t::T
    c::SVector{N,T}
    diverted::Bool
    symmetric::Bool
    sigma_B0::Int
    sigma_Ip::Int
end

"""
    solovev(B0, R0, ϵ, δ, κ, q⋆; B0_dir=1, Ip_dir=1, diverted=false, xpoint=nothing, symmetric = true)

Creates a SolovevEquilibrium Structure

Arguments:\\
`B0` - Toroidal magnetic field magnitude on axis [T]\\
`R0` - Major Radius [m]\\
`ϵ` - inverse aspect ratio a/R0\\
`δ` - triangularity\\
`κ` - elongation/ellipticity\\
`α` - constant relating beta regime (α)\\
`q⋆` - Kink safety factor\\

Keyword Arguments:\\
`B0_dir` - Direction of the toroidal field (+1 = CCW, -1 = CW)\\
`Ip_dir` - Direction of the plasma current (+1 = CCW, -1 = CW)\\
`diverted` - If true then equilibrium has one or more x-points\\
`xpoint` - If diverted = true, then xpoint is set to (R0*(1-1.1*δ*ϵ), -R0*1.1*κ*ϵ) else nothing\\
`symmetric` - Is equilibrium up-down symmetric\\
"""
function solovev(B0, R0, ϵ, δ, κ, α, qstar;
                 B0_dir = 1, Ip_dir = 1,
                 diverted::Bool = false,
                 xpoint::Union{NTuple{2},Nothing} = (diverted ? (R0*(1-1.1*δ*ϵ),-R0*1.1*κ*ϵ) : nothing),
                 symmetric::Bool = (xpoint === nothing))

    if δ > sin(1)
        @warn "Equilibrium is not convex. δ > sin(1)"
    end

    diverted = xpoint !== nothing

    #Eq. 11
    δ₀ = asin(δ)
    N₁ = -(1 + δ₀)^2/(ϵ*κ^2) # [d²x/dy²]_(τ=0)
    N₂ =  (1 - δ₀)^2/(ϵ*κ^2) # [d²x/dy²]_(τ=pi)
    N₃ = -κ/(ϵ*cos(δ₀)^2)    # [d²y/dx²]_(τ=pi/2)

    Ψ_h = zeros(typeof(α),12,12)
    Ψ_p = zeros(typeof(α),12)
    # Boundary Conditions: (ψ_p - bc) + ψ_h'*c = 0
    if symmetric
        if xpoint === nothing
            # Eq. 10
            # Outer equatorial point
            x, y = 1 + ϵ, 0
            Ψ_p[1] = _solovev_psi_p(x, y, α)
            Ψ_h[:,1] .= _solovev_polynomials(x,y)
            # Inner equatorial point
            x, y = 1 - ϵ, 0
            Ψ_p[2] = _solovev_psi_p(x, y, α)
            Ψ_h[:,2] .= _solovev_polynomials(x, y)
            # High point
            x, y = 1 - δ*ϵ, κ*ϵ
            Ψ_p[3] = _solovev_psi_p(x, y, α)
            Ψ_h[:,3] .= _solovev_polynomials(x, y)
            # High point maximum
            x, y = 1 - δ*ϵ, κ*ϵ
            Ψ_p[4] = _solovev_psi_p_Dx(x, y, α)
            Ψ_h[:,4] .= _solovev_polynomials_Dx(x, y)
            # Outer equatorial point curvature
            x, y = 1 + ϵ, 0
            Ψ_p[5] = _solovev_psi_p_Dyy(x, y, α) + N₁*_solovev_psi_p_Dx(x, y, α)
            Ψ_h[:,5] .= _solovev_polynomials_Dyy(x, y) .+ N₁*_solovev_polynomials_Dx(x, y)
            # Inner equatorial point curvature
            x, y = 1 - ϵ, 0
            Ψ_p[6] = _solovev_psi_p_Dyy(x, y, α) + N₂*_solovev_psi_p_Dx(x, y, α)
            Ψ_h[:,6] .= _solovev_polynomials_Dyy(x, y) .+ N₂*_solovev_polynomials_Dx(x, y)
            # High point curvature
            x, y = 1 - δ*ϵ, κ*ϵ
            Ψ_p[7] = _solovev_psi_p_Dxx(x, y, α) + N₃*_solovev_psi_p_Dy(x, y, α)
            Ψ_h[:,7] .= _solovev_polynomials_Dxx(x, y) .+ N₃*_solovev_polynomials_Dy(x, y)

            c = SVector{7}(Ψ_h[1:7,1:7]'\(-Ψ_p[1:7]))
        else
            xsep, ysep = xpoint[1]/R0, xpoint[2]/R0 # normalize xpoint

            # Eq. 12
            # Outer equatorial point
            x, y = 1 + ϵ, 0
            Ψ_p[1] = _solovev_psi_p(x, y, α)
            Ψ_h[:,1] .= _solovev_polynomials(x,y)
            # Inner equatorial point
            x, y = 1 - ϵ, 0
            Ψ_p[2] = _solovev_psi_p(x, y, α)
            Ψ_h[:,2] .= _solovev_polynomials(x, y)
            # High point
            x, y = xsep, ysep
            Ψ_p[3] = _solovev_psi_p(x, y, α)
            Ψ_h[:,3] .= _solovev_polynomials(x, y)
            # B_norm = 0 at high point
            x, y = xsep, ysep
            Ψ_p[4] = _solovev_psi_p_Dx(x, y, α)
            Ψ_h[:,4] .= _solovev_polynomials_Dx(x, y)
            # B_tang = 0 at high point
            x, y = xsep, ysep
            Ψ_p[5] = _solovev_psi_p_Dy(x, y, α)
            Ψ_h[:,5] .= _solovev_polynomials_Dy(x, y)
            # Outer equatorial point curvature
            x, y = 1 + ϵ, 0
            Ψ_p[6] = _solovev_psi_p_Dyy(x, y, α) + N₁*_solovev_psi_p_Dx(x, y, α)
            Ψ_h[:,6] .= _solovev_polynomials_Dyy(x, y) .+ N₁*_solovev_polynomials_Dx(x, y)
            # Inner equatorial point curvature
            x, y = 1 - ϵ, 0
            Ψ_p[7] = _solovev_psi_p_Dyy(x, y, α) + N₂*_solovev_psi_p_Dx(x, y, α)
            Ψ_h[:,7] .= _solovev_polynomials_Dyy(x, y) .+ N₂*_solovev_polynomials_Dx(x, y)

            c = SVector{7}(Ψ_h[1:7,1:7]'\(-Ψ_p[1:7]))
        end
    else
        xsep, ysep = xpoint[1]/R0, xpoint[2]/R0 # normalize xpoint
        if ysep > 0
            throw(ArgumentError("X-point should be below the midplane"))
        end

        # Eq. 28
        # Outer equatorial point
        x, y = 1 + ϵ, 0
        Ψ_p[1] = _solovev_psi_p(x, y, α)
        Ψ_h[:,1] .= _solovev_polynomials(x,y)
        # Inner equatorial point
        x, y = 1 - ϵ, 0
        Ψ_p[2] = _solovev_psi_p(x, y, α)
        Ψ_h[:,2] .= _solovev_polynomials(x, y)
        # Upper high point
        x, y = 1 - δ*ϵ, κ*ϵ
        Ψ_p[3] = _solovev_psi_p(x, y, α)
        Ψ_h[:,3] .= _solovev_polynomials(x, y)
        # Lower X-point
        x, y = xsep, ysep
        Ψ_p[4] = _solovev_psi_p(x, y, α)
        Ψ_h[:,4] .= _solovev_polynomials(x, y)
        # Outer equatorial point up-down symmetry
        x, y = 1 + ϵ, 0
        Ψ_p[5] = _solovev_psi_p_Dy(x,y, α)
        Ψ_h[:,5] .= _solovev_polynomials_Dy(x,y)
        # Inner equatorial point up-down symmetry
        x, y = 1 - ϵ, 0
        Ψ_p[6] = _solovev_psi_p_Dy(x, y, α)
        Ψ_h[:,6] .= _solovev_polynomials_Dy(x, y)
        # Upper high point maximum
        x, y = 1 - δ*ϵ, κ*ϵ
        Ψ_p[7] = _solovev_psi_p_Dx(x, y, α)
        Ψ_h[:,7] .= _solovev_polynomials_Dx(x, y)
        # B_y = 0 at X-point
        x, y = xsep, ysep
        Ψ_p[8] = _solovev_psi_p_Dx(x, y, α)
        Ψ_h[:,8] .= _solovev_polynomials_Dx(x, y)
        # B_x = 0 at X-point
        x, y = xsep, ysep
        Ψ_p[9] = _solovev_psi_p_Dy(x, y, α)
        Ψ_h[:,9] .= _solovev_polynomials_Dy(x, y)
        # Outer equatorial point curvature
        x, y = 1 + ϵ, 0
        Ψ_p[10] = _solovev_psi_p_Dyy(x, y, α) + N₁*_solovev_psi_p_Dx(x, y, α)
        Ψ_h[:,10] .= _solovev_polynomials_Dyy(x, y) .+ N₁*_solovev_polynomials_Dx(x, y)
        # Inner equatorial point curvature
        x, y = 1 - ϵ, 0
        Ψ_p[11] = _solovev_psi_p_Dyy(x, y, α) + N₂*_solovev_psi_p_Dx(x, y, α)
        Ψ_h[:,11] .= _solovev_polynomials_Dyy(x, y) .+ N₂*_solovev_polynomials_Dx(x, y)
        # High point curvature
        x, y = 1 - δ*ϵ, κ*ϵ
        Ψ_p[12] = _solovev_psi_p_Dxx(x, y, α) + N₃*_solovev_psi_p_Dy(x, y, α)
        Ψ_h[:,12] .= _solovev_polynomials_Dxx(x, y) .+ N₃*_solovev_polynomials_Dy(x, y)

        c = SVector{12}(Ψ_h'\(-Ψ_p))
    end

    τ = range(0,2pi,length=100)
    x = (1 .+ ϵ .* cos.(τ .+ δ₀*sin.(τ)))
    y = (ϵ*κ*sin.(τ))
    bdry = PlasmaBoundary(collect(zip(x,y)))

    #Eq. 6.158 in Ideal MHD
    V = volume(bdry,dx=0.005,dy=0.005)
    K₁ = area_average(bdry, (x,y) -> (α + (1-α)*x^2)/x,dx=0.005,dy=0.005)*area(bdry)
    K₂ = volume_average(bdry,(x,y) -> -_solovev_psi(x,y,α,c),dx=0.005,dy=0.005)*V/(2pi)
    K₃ = V/(2pi)

    #Eq. 19 from paper (error in β_t in Ideal MHD (extra ^2))
    Cp = circumference(bdry)
    a = ϵ*R0
    ψ0 = -Cp*a*R0*B0/(qstar*K₁*Ip_dir)
    β_p = 2*(1-α)*Cp^2 * K₂/(K₁^2*K₃)
    β_t = ϵ^2*β_p/(qstar^2)

    cc = 3
    return SolovevEquilibrium(cocos(cc),promote(B0,R0,ϵ,δ,κ,α,qstar,ψ0,β_p,β_t)...,c, diverted, symmetric, B0_dir, Ip_dir)
end

function (S::SolovevEquilibrium)(r,z)
    x = r/S.R0
    y = z/S.R0
    logx = log(x)
    return S.psi0*_solovev_psi(x,y,S.alpha,S.c,logx=logx)
end

function Base.show(io::IO, S::SolovevEquilibrium)
    print(io, "SolovevEquilibrium\n")
    print(io, "  B0 = $(S.sigma_B0*S.B0) [T]\n")
    print(io, "  R0 = $(S.R0) [m]\n")
    print(io, "  ϵ  = $(S.epsilon)\n")
    print(io, "  δ  = $(S.delta)\n")
    print(io, "  κ  = $(S.kappa)\n")
    print(io, "  α  = $(S.alpha)\n")
    print(io, "  q⋆ = $(S.qstar)\n")
    print(io, "  βp = $(S.beta_p)\n")
    print(io, "  βt = $(S.beta_t)\n")
    print(io, "  σ  = $(S.sigma_B0*S.sigma_Ip)\n")
    print(io, "  diverted  = $(S.diverted)\n")
    print(io, "  symmetric = $(S.symmetric)")
end

Base.broadcastable(S::SolovevEquilibrium) = (S,)

function cocos(S::SolovevEquilibrium)
    return S.cocos
end

function B0Ip_sign(S::SolovevEquilibrium)
    return S.sigma_B0*S.sigma_Ip
end

function limits(S::SolovevEquilibrium)
    xlims = (S.R0*(1-1.2*S.epsilon), S.R0*(1 + 1.2*S.epsilon))
    ylims = (S.R0*(-1.2*S.epsilon*S.kappa), S.R0*(1.2*S.epsilon*S.kappa))
    return xlims, ylims
end

function boundary(S::SolovevEquilibrium; n=100, kwargs...)
    if S.diverted
        return boundary(S,0.0; kwargs...)
    else
        τ = range(0,2pi,length=n)
        δ₀ = asin(S.delta)
        x = S.R0*(1 .+ S.epsilon .* cos.(τ .+ δ₀*sin.(τ)))
        y = S.R0*(S.epsilon*S.kappa*sin.(τ))
        return PlasmaBoundary(collect(zip(x,y)))
    end
end

function psi_gradient(S::SolovevEquilibrium,r,z)
    x, y = r/S.R0, z/S.R0
    logx=log(x)
    return (S.psi0/S.R0)*SVector{2}(_solovev_psi_Dx(x,y,S.alpha,S.c; logx=logx), _solovev_psi_Dy(x,y,S.alpha,S.c;logx=logx))
end

_solovev_magnetic_axis = Dict{SolovevEquilibrium,NTuple{2}}()
function clear_cache(S)
    delete!(_solovev_magnetic_axis,S)
end

function magnetic_axis(S::SolovevEquilibrium; x0 = (S.R0, zero(S.R0)))
    if S in keys(_solovev_magnetic_axis)
        return _solovev_magnetic_axis[S]
    else
        sigma_psi = sign(S(x0[1],x0[2]))
        res = Optim.optimize(x->-sigma_psi*S(x[1],x[2]), x -> -sigma_psi*psi_gradient(S,x[1],x[2]), collect(x0), inplace=false)
        axis = (res.minimizer[1],res.minimizer[2])
        _solovev_magnetic_axis[S] = axis
    end
    return axis
end

function psi_limits(S::SolovevEquilibrium)
    psimag = S(magnetic_axis(S)...)
    psibry = zero(psimag)
    return (psimag, psibry)
end

function pressure_gradient(S::SolovevEquilibrium)
    C = S.psi0*(1-S.alpha)/(S.R0^4)
    return -C/mu0
end

function pressure_gradient(S::SolovevEquilibrium, psi)
    return pressure_gradient(S) + psi*0
end

function pressure(S::SolovevEquilibrium, psi; p0=zero(psi))
    C = S.psi0*(1-S.alpha)/(S.R0^4)
    dpdpsi = -C/mu0
    p = dpdpsi*psi + p0
    return p
end

function poloidal_current(S::SolovevEquilibrium,psi)
    A = S.psi0*S.alpha/(S.R0^2)
    dF2dpsi = -2*A
    F2 = abs(dF2dpsi*psi + S.R0^2*S.B0^2)
    return S.sigma_B0*sqrt(F2)
end

function poloidal_current_gradient(S::SolovevEquilibrium,psi)
    A = S.psi0*S.alpha/(S.R0^2)
    return -A/poloidal_current(S,psi)
end

electric_potential(S::SolovevEquilibrium,psi) = zero(psi)
electric_potential_gradient(S::SolovevEquilibrium,psi) = zero(psi)

# === Special Cases ===

function curlB(S::SolovevEquilibrium, r, z)
    # curlB = mu0*J
    psi = S(r,z)
    gval = poloidal_current(S,psi)
    grad_psi = psi_gradient(S,r,z)

    gp = poloidal_current_gradient(S,psi)
    pp = pressure_gradient(S)

    cc = cocos(S)

    invr = inv(r)
    cocos_factor = cc.sigma_RpZ*invr
    cBr = -cocos_factor*gp*grad_psi[2]
    cBz =  cocos_factor*gp*grad_psi[1]

    cocos_factor = -cc.sigma_Bp*((2pi)^cc.exp_Bp)
    cBt = cocos_factor*(r*pp*mu0 + gval*gp*invr)

    return SVector{3}(cylindrical_cocos(cc,cBr,cBt,cBz))
end

function curlB(S::SolovevEquilibrium, x, y, z)
    r = hypot(x,y)
    phi = atan(y,x)
    cB = curlB(S,r,z)
    sp, cp = sincos(phi)
    cB_xyz = SVector{3}(cB[1]*cp - cB[2]*sp, cB[1]*sp + cB[2]*cp, cB[3])
    return cB_xyz
end


module MeshSampling

using StaticArrays
using GeometryTypes
using FileIO

export samplePointsFromMesh

function getFaceArea(points, face::Face)
    a = points[face[1]]
    b = points[face[2]]
    c = points[face[3]]
    return norm(cross(b - a, c - a)) / 2.
end

function samplePointsFromMesh(N_points, points, faces)
    # Get normalized cumulative area of faces
    cum_area = Array{Float64}(size(faces, 1))
    for i = 1:size(faces, 1)
        cum_area[i] = getFaceArea(points, faces[i])
    end
    cum_area = cumsum(cum_area)
    cum_area /= cum_area[end]

    sampled_pts = Array{Float64}(3, N_points)
    sampled_faces = zeros(N_points, size(faces, 1))
    for i=1:N_points
        # Pick the face we'll sample from
        sample = rand()
        k = findfirst(cum_area .>= sample)
        a = points[faces[k][1]]
        b = points[faces[k][2]]
        c = points[faces[k][3]]

        sampled_faces[i, k] = 1

        s1 = rand()
        s2 = rand()
        while (s1 + s2 > 1.0)
            s1 = rand()
            s2 = rand()
        end
        pt = a + s1*(b - a) + s2*(c - a)
        for k=1:3
            sampled_pts[k, i] = pt[k]
        end
    end
    return (sampled_pts, sampled_faces)
end

end # module


# # Lorenz equations
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/models/Lorenz.ipynb)
#
#md # !!! note "Overview"
#md #     System type: Continuous blackbox system\
#md #     State dimension: 3\
#md #     Application domain: Atmospheric convection
#
# ## Model description
#
# The model is a system of three ordinary differential equations now known as
# the Lorenz equations:
#
# ```math
#   \begin{array}{lcl}
#   \dfrac{dx}{dt} & = & \sigma (y-x), \\ \\
#   \dfrac{dy}{dt} & = & x(\rho - z) - y, \\ \\
#   \dfrac{dz}{dt} & = & xy-\beta z
#   \end{array}
# ```
#
# The equations relate the properties of a two-dimensional fluid layer uniformly
# warmed from below and cooled from above. In particular, the equations describe
# the rate of change of three quantities with respect to time:``x`` is
# proportional to the rate of convection, ``y`` to the horizontal temperature
# variation, and ``z`` to the vertical temperature variation. The constants
# ``\sigma`` , ``\rho`` , and ``\beta``  are system parameters proportional to
# the Prandtl number, Rayleigh number, and certain physical dimensions of the
# layer itself.

using ReachabilityAnalysis, Plots

@taylorize function lorenz!(dx, x, params, t)
    local σ = 10.0
    local β = 8.0 / 3.0
    local ρ = 28.0
    dx[1] = σ * (x[2] - x[1])
    dx[2] = x[1] * (ρ - x[3]) - x[2]
    dx[3] = x[1] * x[2] - β * x[3]
    return dx
end

# ## Reachability settings
#
# The initial values considered are ``X_0 \in [0.9, 1.1] \times [0, 0] \times [0, 0]``,
# for a time span of `10`.

X0 = Hyperrectangle(low=[0.9, 0.0, 0.0], high=[1.1, 0.0, 0.0])
prob = @ivp(x' = lorenz!(x), dim=3, x(0) ∈ X0);

# ## Results

# We compute the flowpipe using the TMJets algorithm with ``n_T=10`` and ``n_Q=2``.

alg = TMJets(abstol=1e-15, orderT=10, orderQ=2, maxsteps=50_000);

sol = solve(prob, T=10.0, alg=alg);

solz = overapproximate(sol, Zonotope);

#-

plot(solz, vars=(0, 1), xlab="t", ylab="x")

# It is apparent by inspection that variable $x(t)$ does not exceed 20.0 in the computed
# time span:

plot(solz(0.0 .. 1.5), vars=(0, 1), xlab="t", ylab="x", lw=0.0)
plot!(x -> 20.0, c=:red, xlims=(0.0, 1.5), lab="")

# We can prove that it is the case by computing the support function of the flowpipe
# along direction `[1.0, 0.0, 0.0]`:

ρ([1.0, 0.0, 0.0], solz(0.0 .. 1.5))

# In a similar fashion, we can compute extremal values of variable $y(t)$:

plot(solz, vars=(0, 2), xlab="t", ylab="y")

# Since we have computed overapproximations of the exact flowipe, the following
# quantities are a lower bound on the exact minimum (resp. an uppper bound on the
# exact maximum):

-ρ([0.0, -1.0, 0.0], solz)

ρ([0.0, 1.0, 0.0], solz)

#-

plot(solz, vars=(0, 3), xlab="t", ylab="z")

#-

# Below we plot the flowpipe projected on the `(1, 3)` plane.

plot(solz, vars=(1, 3))


#=

Pkg.add("JLD2")
Pkg.add("Random")
Pkg.add("LinearAlgebra")
Pkg.add("Statistics")
Pkg.add("CSV")
Pkg.add("DataFrames")
Pkg.add("FreqTables")

Pkg.add("Kronecker")

=#

import JLD2
import Random
import LinearAlgebra
import Statistics
import CSV
import DataFrames
import FreqTables

import Kronecker

###################################################################
### Question 1

#function q1()

## (a)

Random.seed!(1234)

# i matrix 10* 7 with uniform values between -5 and 10

A=rand(-5:10,(10,7)) # -5:10 is the range, 10,7 is the size of the matrix

# ii matrix 10*7 with normal values (-2,15) => not re-centered, do not know how

B=randn(10,7)*15 # 15 is the standard deviation

# iii put 5 rows * 5 columns of the first together with 5 rows * 2 columns (last) of the second

A_small=A[1:5,:1:5]
A_small
B_small=B[1:5,6:7]
B_small

C=[A_small B_small] # same on the row dimension only 5*5 and 5*2 matrix

size(C)

# iv 10*7 matrix where values are equal to As if positive, else 0s

A

D=A

A # it seems that this changes M1 too - Why?

for row=1:10
    for column=1:7  # goes through all arguments of the matrix
        if D[row,column]<0
            D[row,column]=0
        end
    end
end

D

## (b) get the number of elements of A

length(A)

## (c) list the number of unique elements in D

nr_unique=length(unique(D))

# (d)

println(reshape(B,(70,1))) # reshaped it into a vector 70*1

size(reshape(B,(70,1))) # 70,1

size(reshape(B,(1,70))) # 1,70

# (e)

E=rand(2,10,7) # a matrix with the right dimensions but random numbers

E[1,1:10,1:7]=A

E[2,1:10,1:7]=B

E

# (f)

E=permutedims(E, [2,3,1]) # what used to be the second dimension now is first ...

E

# (g)

F = A ⊗ B

#F = C ⊗ E # no method matching error - Why?

size(A) # 10,7
size(B) # 10,7
size(C) # 5,7
size(E) # 10,7,2 => dimensions do not match

# (h)

using JLD2
@save("C:/Users/tsankova/Dropbox/Structural Course Material/Week 1/matrixpractice.jld", "arr", A, B, C, D, E, F)

# (i)

using JLD2
@save("C:/Users/tsankova/Dropbox/Structural Course Material/Week 1/firstmatrix.jld", "arr", A, B, C, D)

# (j)

using JLD2
@save("C:/Users/tsankova/Dropbox/Structural Course Material/Week 1/Cmatrix.csv", "arr", C)

# (k)

using DataFrames
D=convert(DataFrame, D)

using JLD2
@save("C:/Users/tsankova/Dropbox/Structural Course Material/Week 1/Dmatrix.dat", D)

# (l)

#end

#q1()

###################################################################
### Question 2

#function q2()

# (a)

AB=zeros(10,7)

for row=1:10
    for column=1:7  # goes through all arguments of the matrix
        a=A[row,column]
        println(a)
        b=B[row,column]
        println(b)
        c=a*b
        println(c)
        AB[row,column]=c
    end
end

# (b)

C

Cprime=reshape(C,(35,1))

Cprime=Vector{Float64}(vec(Cprime))

Cprime

try
    for row=1:length(Cprime) # length changes as drop obs. so in the end there is an error message
        if Cprime[row]>5
            deleteat!(Cprime,row)
        elseif Cprime[row]<-5
            deleteat!(Cprime,row)
        end
    end
catch
finally
end

Cprime # 26 elements left

# Do it not in a loop => left for later

# (c)

# (d)

# (e)

# (f)

#end

#q2()

###################################################################
### Question 3

#function q3()

# (a)

using CSV
nlsw88 = CSV.read("nlsw88.csv")

names(nlsw88) # lists variable names, seem file

using Statistics
mean(nlsw88.grade) # mean of grade variable - missing?

# missing values?

using JLD2
@save("C:/Users/tsankova/Dropbox/Structural Course Material/Week 1/nlsw88.jld", "df", nlsw88)

# (b)

# share married - married variable

using Statistics
mean(nlsw88.married) # 0.642

# share with a college degree - collgrad variable

using Statistics
mean(nlsw88.collgrad) # 0.237

# (c) do not know how to show percentag ???

using FreqTables
freqtable(nlsw88.race)

# (d)

summarystats=describe(nlsw88) # min, max, median, mean values

nlsw88

using FreqTables
freqtable(nlsw88.grade) # 2 missing

# (e)

using FreqTables
freqtable(nlsw88.industry)
freqtable(nlsw88.occupation)
freqtable(nlsw88.industry,nlsw88.occupation)

# (f)

ind_wage=select(nlsw88,:industry,:wage) # select data wanted
ind_wage_grouped=groupby(ind_wage,1) # create a grouped data frame
ind_wage_selected=combine(ind_wage_grouped,:wage => mean) # calculate mean wage by group

occ_wage=select(nlsw88,:occupation,:wage)
occ_wage_grouped=groupby(occ_wage,1)
occ_wage_selected=combine(occ_wage_grouped,:wage => mean)

#end

#q3()

###################################################################
### Question 4

#function q4()

# (a)

@load "firstmatrix.jld"

# (b) and (c)

function matrixops(A_some,B_some)

    #=

    The function takes A and B matrices as agruments of the same
    dimension, then creates three matrices - one element-by-element
    multiplication, one matrix multiplication and one matrix
    summation

    =#

    size_A=size(A_some)

    size_A_r=size_A[1] # number of rows
    size_A_c=size_A[2] # number of columns

    e_by_e_product_A_B=zeros(size_A_r,size_A_c)

    for row in 1:size_A_r
        for column in 1:size_A_c
            a=A_some[row,column]
            b=B_some[row,column]
            c=a*b
            e_by_e_product_A_B[row,column]=c
        end
    end

    A_transpose_B=transpose(A_some)*B_some

    A_plus_B=A_some+B_some

    return e_by_e_product_A_B, A_transpose_B, A_plus_B

end

# (d)

matrixops(A,B) # returns three matrices of the correct size

# (e)

function matrixops(A_some,B_some)

    size_A=size(A_some)
    size_B=size(B_some)

    if size_A==size_B # like stata

        size_A_r=size_A[1] # number of rows
        size_A_c=size_A[2] # number of columns

        e_by_e_product_A_B=zeros(size_A_r,size_A_c)

        for row in 1:size_A_r
            for column in 1:size_A_c
                a=A_some[row,column]
                b=B_some[row,column]
                c=a*b
                e_by_e_product_A_B[row,column]=c
            end
        end

        A_transpose_B=transpose(A_some)*B_some

        A_plus_B=A_some+B_some

        return e_by_e_product_A_B, A_transpose_B, A_plus_B

    else

        println("Matrices of different sizes")

    end

end

# (e)

matrixops(C,D) # prints the message

# (f)

ttl_exp=convert(Array,nlsw88.ttl_exp)

wage=convert(Array,nlsw88.wage)

ttl_exp
wage # vectors

ttl_exp_array=reshape(ttl_exp, length(ttl_exp), 1)
wage_array=reshape(wage, length(wage), 1)

matrixops(ttl_exp_array,wage_array) # yes

#end

#q4()

#=

List of issues:

Do not know how to fix the normal distribution
to have a different mean and standard deviation

Have not done the second part of Question 2 -
need to see answers

Q: Do we write using ...
when the packages have sub-packaged / commands with the same name
only or we should always do it? Seems wordy

Q: Why do we need three dimensional arrays?

Q: If I want to replace the values of a matrix with something
one by one then I first need to create an empty / zeros matrix -
is this true?

Do not know how to declare the functions
without inputs and outpits
- using expressions need to come forward
- some other pieces of code do not run (kron product)
- unclear when to use return as well

Q: Is it good practice to wrap everything in a function?

Q: Should I have se a home directory?


# This file is part of the TaylorIntegration.jl package; MIT licensed

using TaylorIntegration, Elliptic
using LinearAlgebra: norm
using Test
using Logging
import Logging: Warn

@testset "Testing `jettransport.jl`" begin

    local _order = 28
    local _abstol = 1.0E-20

    max_iters_reached() = "Maximum number of integration steps reached; exiting.\n"

    f(x, p, t) = x^2
    g(x, p, t) = 0.3x

    @testset "Test TaylorN jet transport (t0, tmax): 1-dim case" begin
        p = set_variables("ξ", numvars=1, order=5)
        x0 = 3.0 #"nominal" initial condition
        x0TN = x0 + p[1] #jet transport initial condition
        t0=0.0
        tmax=0.3
        tvTN, xvTN = (@test_logs (Warn, max_iters_reached()) taylorinteg(
            f, x0TN, t0, tmax, _order, _abstol, maxsteps=1))
        @test size(tvTN) == (2,)
        @test size(xvTN) == (2,)
        tvTN, xvTN = (@test_logs min_level=Logging.Warn taylorinteg(
            f, x0TN, t0, tmax, _order, _abstol))
        tv, xv = (@test_logs min_level=Logging.Warn taylorinteg(
            f, x0, t0, tmax, _order, _abstol))
        exactsol(t, x0, t0) = x0/(1.0-x0*(t-t0)) #the analytical solution
        δsol = exactsol(tvTN[end], x0TN, t0)-xvTN[end]
        δcoeffs = map(y->y[1], map(x->x.coeffs, δsol.coeffs))
        @test isapprox(δcoeffs, zeros(6), atol=1e-10, rtol=0)
        @test isapprox(x0, evaluate(xvTN[1]))
        @test isapprox(xv[1], evaluate(xvTN[1]))
        @test isapprox(xv[end], evaluate(xvTN[end]))
        xvTN_0 = evaluate.(xvTN)
        @test norm(exactsol.(tvTN, x0, t0)-xvTN_0, Inf) < 1E-12
        for i in 1:5
            disp = 0.001*rand() #a small, random displacement
            x0_disp = x0+disp
            dv = map(x->[disp], tvTN) #a vector of identical displacements
            tv_disp, xv_disp = (@test_logs min_level=Logging.Warn taylorinteg(
                f, x0_disp, t0, tmax, _order, _abstol))
            xvTN_disp = evaluate.(xvTN, dv)
            @test norm(exactsol.(tvTN, x0_disp, t0)-xvTN_disp, Inf) < 1E-12 #analytical vs jet transport
            @test norm(x0_disp-evaluate(xvTN[1], [disp]), Inf) < 1E-12
            @test norm(xv_disp[1]-evaluate(xvTN[1], [disp]), Inf) < 1E-12
            @test norm(xv_disp[end]-evaluate(xvTN[end], [disp]), Inf) < 1E-12
        end

        y0 = 1.0 #"nominal" initial condition
        u0 = 0.0 #initial time
        y0TN = y0 + p[1] #jet transport initial condition
        uvTN, yvTN = (@test_logs (Warn, max_iters_reached()) taylorinteg(
            g, y0TN, u0, 10/0.3, _order, _abstol, maxsteps=1))
        @test size(uvTN) == (2,)
        @test size(yvTN) == (2,)
        uvTN, yvTN = (@test_logs min_level=Logging.Warn taylorinteg(
            g, y0TN, u0, 10/0.3, _order, _abstol))
        uv, yv = (@test_logs min_level=Logging.Warn taylorinteg(
            g, y0, u0, 10/0.3, _order, _abstol))
        exactsol_g(u, y0, u0) = y0*exp(0.3(u-u0))
        δsol_g = exactsol_g(uvTN[end], y0TN, u0)-yvTN[end]
        δcoeffs_g = map(y->y[1], map(x->x.coeffs, δsol_g.coeffs))
        @test isapprox(δcoeffs_g, zeros(6), atol=1e-10, rtol=0)
        @test isapprox(y0, evaluate(yvTN[1]), atol=1e-10, rtol=0)
        @test isapprox(yv[1], evaluate(yvTN[1]), atol=1e-10, rtol=0)
        @test isapprox(yv[end], evaluate(yvTN[end]), atol=1e-10, rtol=0)
        yvTN_0 = evaluate.(yvTN)
        @test norm(exactsol_g.(uvTN, y0, u0)-yvTN_0, Inf) < 1E-9
        for i in 1:5
            disp = 0.001*rand() #a small, random displacement
            y0_disp = y0+disp
            dv = map(x->[disp], uvTN) #a vector of identical displacements
            uv_disp, yv_disp = (@test_logs min_level=Logging.Warn taylorinteg(
                g, y0_disp, u0, 10/0.3, _order, _abstol))
            yvTN_disp = evaluate.(yvTN, dv)
            @test norm(exactsol_g.(uvTN, y0_disp, u0)-yvTN_disp, Inf) < 1E-9 #analytical vs jet transport
            @test norm(y0_disp-yvTN[1]([disp]), Inf) < 1E-9
            @test norm(yv_disp[1]-yvTN[1]([disp]), Inf) < 1E-9
            @test norm(yv_disp[end]-yvTN[end]([disp]), Inf) < 1E-9
        end
    end

    @testset "Test Taylor1 jet transport (t0, tmax): 1-dim case" begin
        p = Taylor1([0.0,1.0], 5)
        x0 = 3.0 #"nominal" initial condition
        x0T1 = x0 + p #jet transport initial condition
        t0 = 0.0
        tmax = 0.3
        tvT1, xvT1 = (@test_logs (Warn, max_iters_reached()) taylorinteg(
            f, x0T1, t0, tmax, _order, _abstol, maxsteps=1))
        @test size(tvT1) == (2,)
        @test size(xvT1) == (2,)
        tvT1, xvT1 = (@test_logs min_level=Logging.Warn taylorinteg(
            f, x0T1, t0, tmax, _order, _abstol))
        tv, xv = (@test_logs min_level=Logging.Warn taylorinteg(
            f, x0, t0, tmax, _order, _abstol))
        exactsol(t, x0, t0) = x0/(1.0-x0*(t-t0)) #the analytical solution
        δsol = exactsol(tvT1[end], x0T1, t0)-xvT1[end] #analytical vs jet transport diff at end of integration
        δcoeffs = δsol.coeffs
        @test isapprox(δcoeffs, zeros(6), atol=1e-10, rtol=0)
        @test isapprox(x0, evaluate(xvT1[1]))
        @test isapprox(xv[1], evaluate(xvT1[1]))
        @test isapprox(xv[end], evaluate(xvT1[end]))
        xvT1_0 = evaluate.(xvT1)
        @test norm(exactsol.(tvT1, x0, t0)-xvT1_0, Inf) < 1E-12
        for i in 1:5
            disp = 0.001*rand() #a small, random displacement
            x0_disp = x0+disp
            tv_disp, xv_disp = (@test_logs min_level=Logging.Warn taylorinteg(
                f, x0_disp, t0, tmax, _order, _abstol))
            xvT1_disp = evaluate.(xvT1, disp)
            @test norm(exactsol.(tvT1, x0_disp, t0)-xvT1_disp, Inf) < 1E-12 #analytical vs jet transport
            @test norm(x0_disp-evaluate(xvT1[1], disp), Inf) < 1E-12
            @test norm(xv_disp[1]-evaluate(xvT1[1], disp), Inf) < 1E-12
            @test norm(xv_disp[end]-evaluate(xvT1[end], disp), Inf) < 1E-12
        end

        y0 = 1.0 #"nominal" initial condition
        u0 = 0.0 #initial time
        y0T1 = y0 + p #jet transport initial condition
        uvT1, yvT1 = (@test_logs (Warn, max_iters_reached()) taylorinteg(
            g, y0T1, t0, 10/0.3, _order, _abstol, maxsteps=1)) #warmup lap
        # test maxsteps break
        @test size(uvT1) == (2,)
        @test size(yvT1) == (2,)
        uvT1, yvT1 = (@test_logs min_level=Logging.Warn taylorinteg(
            g, y0T1, t0, 10/0.3, _order, _abstol)) #Taylor1 jet transport integration
        uv, yv = (@test_logs min_level=Logging.Warn taylorinteg(
            g, y0, t0, 10/0.3, _order, _abstol)) #reference integration
        exactsol_g(u, y0, u0) = y0*exp(0.3(u-u0))  #the analytical solution
        δsol_g = exactsol_g(uvT1[end], y0T1, t0)-yvT1[end] #analytical vs jet transport diff at end of integration
        δcoeffs_g = δsol_g.coeffs
        @test isapprox(δcoeffs_g, zeros(6), atol=1e-10, rtol=0)
        @test isapprox(y0, evaluate(yvT1[1]), atol=1e-10, rtol=0)
        @test isapprox(yv[1], evaluate(yvT1[1]), atol=1e-10, rtol=0)
        @test isapprox(yv[end], evaluate(yvT1[end]), atol=1e-10, rtol=0)
        yvT1_0 = evaluate.(yvT1)
        @test norm(exactsol_g.(uvT1, y0, u0)-yvT1_0, Inf) < 1E-9
        for i in 1:5
            disp = 0.001*rand() #a small, random displacement
            y0_disp = y0+disp
            uv_disp, yv_disp = (@test_logs min_level=Logging.Warn taylorinteg(
                g, y0_disp, u0, 10/0.3, _order, _abstol))
            yvT1_disp = evaluate.(yvT1, disp)
            @test norm(exactsol_g.(uvT1, y0_disp, u0)-yvT1_disp, Inf) < 1E-9 #analytical vs jet transport
            @test norm(y0_disp-evaluate(yvT1[1], disp), Inf) < 1E-9
            @test norm(yv_disp[1]-evaluate(yvT1[1], disp), Inf) < 1E-9
            @test norm(yv_disp[end]-evaluate(yvT1[end], disp), Inf) < 1E-9
        end
    end

    @testset "Test Taylor1 jet transport (trange): 1-dim case" begin
        p = Taylor1([0.0,1.0], 5)
        x0 = 3.0 #"nominal" initial condition
        x0T1 = x0 + p #jet transport initial condition
        tv = 0.0:0.05:0.33
        xvT1 = (@test_logs (Warn, max_iters_reached()) taylorinteg(
            f, x0T1, tv, _order, _abstol, maxsteps=1))
        @test size(xvT1) == (7,)
        ta = vec(tv)
        xvT1 = (@test_logs (Warn, max_iters_reached()) taylorinteg(
            f, x0T1, ta, _order, _abstol, maxsteps=1))
        @test size(xvT1) == (7,)
        xvT1 = (@test_logs min_level=Logging.Warn taylorinteg(
            f, x0T1, tv, _order, _abstol))
        xvT11 = (@test_logs min_level=Logging.Warn taylorinteg(
            f, x0T1, ta, _order, _abstol))
        @test xvT11 == xvT1
        xv = (@test_logs min_level=Logging.Warn taylorinteg(
            f, x0, tv, _order, _abstol))
        exactsol(t, x0, t0) = x0/(1.0-x0*(t-t0)) #the analytical solution
        δsol = exactsol(tv[end], x0T1, tv[1])-xvT1[end]
        δcoeffs = δsol.coeffs
        @test length(tv) == length(xvT1)
        @test isapprox(δcoeffs, zeros(6), atol=1e-10, rtol=0)
        xv_analytical = exactsol.(tv, x0, tv[1])
        xvT1_0 = xvT1.()
        @test isapprox(xv_analytical, xvT1_0, atol=1e-10, rtol=0)
        for i in 1:5
            disp = 0.001*rand() #a small, random displacement
            xv_disp = (@test_logs min_level=Logging.Warn taylorinteg(
                f, x0+disp, tv, _order, _abstol))
            xv_disp2 = (@test_logs min_level=Logging.Warn taylorinteg(
                f, x0+disp, ta, _order, _abstol))
            @test xv_disp == xv_disp2
            xvT1_disp = xvT1.(disp)
            @test norm(exactsol.(tv, x0+disp, tv[1])-xvT1_disp, Inf) < 1E-12 #analytical vs jet transport
            @test norm(xv_disp-xvT1_disp, Inf) < 1E-12 # integration vs jet transport
        end

        y0 = 1.0 #"nominal" initial condition
        y0T1 = y0 + p #jet transport initial condition
        uv = 0.0:1/0.3:10/0.3
        yvT1 = (@test_logs (Warn, max_iters_reached()) taylorinteg(
            g, y0T1, uv, _order, _abstol, maxsteps=1))
        @test size(yvT1) == (11,)
        yvT1 = (@test_logs min_level=Logging.Warn taylorinteg(
            g, y0T1, uv, _order, _abstol))
        yv = (@test_logs min_level=Logging.Warn taylorinteg(
            g, y0, uv, _order, _abstol))
        exactsol_g(u, y0, u0) = y0*exp(0.3(u-u0))
        δsol_g = exactsol_g(uv[end], y0T1, uv[1])-yvT1[end]
        δcoeffs_g = δsol_g.coeffs
        @test length(uv) == length(yvT1)
        @test isapprox(δcoeffs_g, zeros(6), atol=1e-10, rtol=0)
        yv_analytical = exactsol_g.(uv, y0, uv[1])
        yvT1_0 = evaluate.(yvT1)
        @test isapprox(yv_analytical, yvT1_0, atol=1e-10, rtol=0)
        for i in 1:5
            disp = 0.001*rand() #a small, random displacement
            yv_disp = (@test_logs min_level=Logging.Warn taylorinteg(
                g, y0+disp, uv, _order, _abstol))
            yvT1_disp = evaluate.(yvT1, disp)
            @test norm(exactsol_g.(uv, y0+disp, uv[1])-yvT1_disp, Inf) < 1E-9 #analytical vs jet transport
            @test norm(yv_disp-yvT1_disp, Inf) < 1E-9 # integration vs jet transport
        end
    end

    @testset "Test TaylorN jet transport (trange): 1-dim case" begin
        p = set_variables("ξ", numvars=1, order=5)
        x0 = 3.0 #"nominal" initial condition
        x0TN = x0 + p[1] #jet transport initial condition
        tv = 0.0:0.05:0.33
        xvTN = (@test_logs (Warn, max_iters_reached()) taylorinteg(
            f, x0TN, tv, _order, _abstol, maxsteps=1))
        @test size(xvTN) == (7,)
        ta = vec(tv)
        xvTN = (@test_logs (Warn, max_iters_reached()) taylorinteg(
            f, x0TN, ta, _order, _abstol, maxsteps=1))
        @test size(xvTN) == (7,)
        xvTN = (@test_logs min_level=Logging.Warn taylorinteg(
            f, x0TN, tv, _order, _abstol))
        xvTN_ = (@test_logs min_level=Logging.Warn taylorinteg(
            f, x0TN, ta, _order, _abstol))
        @test xvTN == xvTN_
        xv = (@test_logs min_level=Logging.Warn taylorinteg(
            f, x0, tv, _order, _abstol))
        exactsol(t, x0, t0) = x0/(1.0-x0*(t-t0)) #the analytical solution
        δsol = exactsol(tv[end], x0TN, tv[1])-xvTN[end]
        δcoeffs = map(y->y[1], map(x->x.coeffs, δsol.coeffs))
        @test length(tv) == length(xvTN)
        @test isapprox(δcoeffs, zeros(6), atol=1e-10, rtol=0)
        xv_analytical = exactsol.(tv, x0, tv[1])
        xvTN_0 = evaluate.(xvTN)
        @test isapprox(xv_analytical, xvTN_0, atol=1e-10, rtol=0)
        for i in 1:5
            disp = 0.001*rand() #a small, random displacement
            dv = map(x->[disp], tv) #a vector of identical displacements
            xv_disp = (@test_logs min_level=Logging.Warn taylorinteg(
                f, x0+disp, tv, _order, _abstol))
            xv_disp_ = (@test_logs min_level=Logging.Warn taylorinteg(
                f, x0+disp, ta, _order, _abstol))
            @test xv_disp == xv_disp_
            xvTN_disp = evaluate.(xvTN, dv)
            @test norm(exactsol.(tv, x0+disp, tv[1])-xvTN_disp, Inf) < 1E-12 #analytical vs jet transport
            @test norm(xv_disp-xvTN_disp, Inf) < 1E-12 # integration vs jet transport
        end

        y0 = 1.0 #"nominal" initial condition
        y0TN = y0 + p[1] #jet transport initial condition
        uv = 0.0:1/0.3:10/0.3
        yvTN = (@test_logs (Warn, max_iters_reached()) taylorinteg(
            g, y0TN, uv, _order, _abstol, maxsteps=1))
        @test size(yvTN) == (11,)
        yvTN = (@test_logs min_level=Logging.Warn taylorinteg(
            g, y0TN, uv, _order, _abstol))
        yv = (@test_logs min_level=Logging.Warn taylorinteg(
            g, y0, uv, _order, _abstol))
        exactsol_g(u, y0, u0) = y0*exp(0.3(u-u0))
        δsol_g = exactsol_g(uv[end], y0TN, uv[1])-yvTN[end]
        δcoeffs_g = map(y->y[1], map(x->x.coeffs, δsol_g.coeffs))
        @test length(uv) == length(yvTN)
        @test isapprox(δcoeffs_g, zeros(6), atol=1e-10, rtol=0)
        yv_analytical = exactsol_g.(uv, y0, uv[1])
        yvTN_0 = evaluate.(yvTN)
        @test isapprox(yv_analytical, yvTN_0, atol=1e-10, rtol=0)
        for i in 1:5
            disp = 0.001*rand() #a small, random displacement
            dv = map(x->[disp], uv) #a vector of identical displacements
            yv_disp = (@test_logs min_level=Logging.Warn taylorinteg(
                g, y0+disp, uv, _order, _abstol))
            yvTN_disp = evaluate.(yvTN, dv)
            @test norm(exactsol_g.(uv, y0+disp, uv[1])-yvTN_disp, Inf) < 1E-9 #analytical vs jet transport
            @test norm(yv_disp-yvTN_disp, Inf) < 1E-9 # integration vs jet transport
        end
    end

    let
        function harmosc!(dx, x, p, t) #the harmonic oscillator ODE
            dx[1] = x[2]
            dx[2] = -x[1]*x[3]^2
            dx[3] = zero(x[1])
            nothing
        end

        @testset "Test Taylor1 jet transport (t0,tmax): harmonic oscillator" begin
            t = Taylor1([0.0, 1.0], 10)
            ω0 = 1.0
            x0 = [0.0,ω0,ω0]
            x0T1 = x0+[0t,t,t]
            tv1, xv1 = (@test_logs (Warn, max_iters_reached()) taylorinteg(
                harmosc!, x0T1, 0.0, 100pi, _order, _abstol, maxsteps=1))
            @test size(tv1) == (2,)
            @test size(xv1) == (2, 3)
            tv1, xv1 = (@test_logs min_level=Logging.Warn taylorinteg(
                harmosc!, x0T1, 0.0, 100pi, _order, _abstol, maxsteps=2000))
            y0 = evaluate.(xv1)
            x1(t,δω) = sin((ω0+δω)*t)
            x2(t,δω) = (ω0+δω)*cos((ω0+δω)*t)
            @test norm(y0[:,1]-x1.(tv1,0.0), Inf) < 1E-11
            @test norm(y0[:,2]-x2.(tv1,0.0), Inf) < 1E-11
            for i in 1:5
                δω=0.001*rand()
                x0_disp = x0+[0.0,δω,δω]
                tv, xv = (@test_logs min_level=Logging.Warn taylorinteg(
                    harmosc!, x0_disp, 0.0, 100pi, _order, _abstol, maxsteps=2000))
                y1 = evaluate.(xv1, δω)
                @test norm(y1[:,1]-x1.(tv1,δω),Inf) < 1E-11
                @test norm(y1[:,2]-x2.(tv1,δω),Inf) < 1E-11
                @test norm(x0_disp-evaluate.(xv1[1,:], δω), Inf) < 1E-11
                @test norm(xv[1,:]-evaluate.(xv1[1,:], δω), Inf) < 1E-11
                @test norm(xv[end,:]-evaluate.(xv1[end,:], δω), Inf) < 1E-11
            end
        end

        @testset "Test Taylor1 jet transport (trange): harmonic oscillator" begin
            t = Taylor1([0.0, 1.0], 10)
            ω0 = 1.0
            x0 = [0.0,ω0,ω0]
            x0T1 = x0+[0t,t,t]
            tv = 0.0:0.25*(2pi):100pi
            xv1 = (@test_logs (Warn, max_iters_reached()) taylorinteg(
                harmosc!, x0T1, tv, _order, _abstol, maxsteps=1))
            @test length(tv) == 201
            @test size(xv1) == (201, 3)
            ta = vec(tv)
            xv1 = (@test_logs (Warn, max_iters_reached()) taylorinteg(
                harmosc!, x0T1, ta, _order, _abstol, maxsteps=1))
            @test length(ta) == 201
            @test size(xv1) == (201, 3)
            xv1 = (@test_logs min_level=Logging.Warn taylorinteg(
                harmosc!, x0T1, tv, _order, _abstol, maxsteps=2000))
            xv1_ = (@test_logs min_level=Logging.Warn taylorinteg(
                harmosc!, x0T1, ta, _order, _abstol, maxsteps=2000))
            @test xv1 == xv1_
            y0 = evaluate.(xv1)
            x1(t,δω) = sin((ω0+δω)*t)
            x2(t,δω) = (ω0+δω)*cos((ω0+δω)*t)
            @test norm(y0[:,1]-x1.(tv,0.0), Inf) < 1E-11
            @test norm(y0[:,2]-x2.(tv,0.0), Inf) < 1E-11
            for i in 1:5
                δω=0.001*rand()
                x0_disp = x0+[0.0,δω,δω]
                xv = (@test_logs min_level=Logging.Warn taylorinteg(
                    harmosc!, x0_disp, tv, _order, _abstol, maxsteps=2000))
                xv_ = (@test_logs min_level=Logging.Warn taylorinteg(
                    harmosc!, x0_disp, ta, _order, _abstol, maxsteps=2000))
                @test xv == xv_
                y1 = evaluate.(xv1, δω)
                @test norm(y1[:,1]-x1.(tv,δω), Inf) < 1E-11
                @test norm(y1[:,2]-x2.(tv,δω), Inf) < 1E-11
                @test norm(y1-xv, Inf) < 1E-11
            end
        end
    end

    let
        function harmosc!(dx, x, p, t) #the harmonic oscillator ODE
            dx[1] = x[2]
            dx[2] = -x[1]
            nothing
        end

        @testset "Test TaylorN jet transport (t0,tmax): harmonic oscillator" begin
            p = set_variables("ξ", numvars=2, order=5)
            x0 = [-1.0,0.45]
            x0TN = x0 + p
            tvTN, xvTN = (@test_logs (Warn, max_iters_reached()) taylorinteg(
                harmosc!, x0TN, 0.0, 10pi, _order, _abstol, maxsteps=1))
            @test length(tvTN) == 2
            @test size(xvTN) == (length(tvTN), length(x0TN))
            tvTN, xvTN = (@test_logs min_level=Logging.Warn taylorinteg(
                harmosc!, x0TN, 0.0, 10pi, _order, _abstol))
            tv , xv = (@test_logs min_level=Logging.Warn taylorinteg(
                harmosc!, x0  , 0.0, 10pi, _order, _abstol))
            x_analyticsol(t,x0,p0) = p0*sin(t)+x0*cos(t)
            p_analyticsol(t,x0,p0) = p0*cos(t)-x0*sin(t)
            x_δsol = x_analyticsol(tvTN[end], x0TN[1], x0TN[2])-xvTN[end,1]
            x_δcoeffs = map(y->y[1], map(x->x.coeffs, x_δsol.coeffs))
            p_δsol = p_analyticsol(tvTN[end], x0TN[1], x0TN[2])-xvTN[end,2]
            p_δcoeffs = map(y->y[1], map(x->x.coeffs, p_δsol.coeffs))

            @test (length(tvTN), length(x0)) == size(xvTN)
            @test isapprox(x_δcoeffs, zeros(6), atol=1e-10, rtol=0)
            @test isapprox(x0, map( x->evaluate(x), xvTN[1,:]))
            @test isapprox(xv[1,:], map( x->evaluate(x), xvTN[1,:])) # nominal solution must coincide with jet evaluated at ξ=(0,0) at initial time
            @test isapprox(xv[end,:], map( x->evaluate(x), xvTN[end,:])) #nominal solution must coincide with jet evaluated at ξ=(0,0) at final time

            xvTN_0 = map( x->evaluate(x), xvTN ) # the jet evaluated at the nominal solution
            @test isapprox(xv[end,:], xvTN_0[end,:]) # nominal solution must coincide with jet evaluated at ξ=(0,0) at final time
            @test isapprox(xvTN_0[1,:], xvTN_0[end,:]) # end point must coincide with a full period
        end
    end

    function pendulum!(dx, x, p, t) #the simple pendulum ODE
        dx[1] = x[2]
        dx[2] = -sin(x[1])
        nothing
    end

    @testset "Test TaylorN jet transport (trange): simple pendulum" begin
        varorder = 2 #the order of the variational expansion
        p = set_variables("ξ", numvars=2, order=varorder) #TaylorN steup
        q0 = [1.3, 0.0] #the initial conditions
        q0TN = q0 + p #parametrization of a small neighbourhood around the initial conditions
        # T is the librational period == 4Elliptic.K(sin(q0[1]/2)^2)
        T = 4Elliptic.K(sin(q0[1]/2)^2) # equals 7.019250311844546
        t0 = 0.0 #the initial time
        tmax = T #the final time
        integstep = 0.25*T #the time interval between successive evaluations of the solution vector

        #the time range
        tr = t0:integstep:tmax;
        #xv is the solution vector representing the propagation of the initial condition q0 propagated until time T
        xv = (@test_logs min_level=Logging.Warn taylorinteg(
            pendulum!, q0, tr, _order, _abstol, maxsteps=100))
        #xvTN is the solution vector representing the propagation of the initial condition q0 plus variations (ξ₁,ξ₂) propagated until time T
        #note that q0 is a Vector{Float64}, but q0TN is a Vector{TaylorN{Float64}}
        #but apart from that difference, we're calling `taylorinteg` essentially with the same parameters!
        #thus, jet transport is reduced to a beautiful application of Julia's multiple dispatch!
        xvTN = (@test_logs (Warn, max_iters_reached()) taylorinteg(
            pendulum!, q0TN, tr, _order, _abstol, maxsteps=1))
        @test size(xvTN) == (5,2)
        xvTN = (@test_logs min_level=Logging.Warn taylorinteg(
            pendulum!, q0TN, tr, _order, _abstol, maxsteps=100))

        xvTN_0 = map( x->evaluate(x, [0.0, 0.0]), xvTN ) # the jet evaluated at the nominal solution

        @test isapprox(xvTN_0[1,:], xvTN_0[end,:]) #end point must coincide with a full period
        @test isapprox(xv, xvTN_0) #nominal solution must coincide with jet evaluated at ξ=(0,0)

        #testing another jet transport method:
        tv, xv = (@test_logs min_level=Logging.Warn taylorinteg(
            pendulum!, q0, t0, tmax, _order, _abstol, maxsteps=100))
        @test isapprox(xv[1,:], xvTN_0[1,:]) #nominal solution must coincide with jet evaluated at ξ=(0,0) at initial time
        @test isapprox(xv[end,:], xvTN_0[end,:]) #nominal solution must coincide with jet evaluated at ξ=(0,0) at final time

        # a small displacement
        disp = 0.0001 #works even with 0.001, but we're giving it some margin

        #compare the jet solution evaluated at various variation vectors ξ, wrt to full solution, at each time evaluation point
        xv_disp = (@test_logs (Warn, max_iters_reached()) taylorinteg(
            pendulum!, q0+[disp,0.0], tr, _order, _abstol, maxsteps=1))
        @test size(xv_disp) == (5,2)
        for i in 1:10
            # generate a random angle
            ϕ = 2pi*rand()
            # generate a random point on a circumference of radius disp
            ξ = disp*[cos(ϕ), sin(ϕ)]
            #propagate in time full solution with initial condition q0+ξ
            xv_disp = (@test_logs min_level=Logging.Warn taylorinteg(
                pendulum!, q0+ξ, tr, _order, _abstol, maxsteps=100))
            #evaluate jet at q0+ξ
            xvTN_disp = map( x->evaluate(x, ξ), xvTN )
            #the propagated exact solution at q0+ξ should be approximately equal to the propagated jet solution evaluated at the same point:
            @test isapprox( xvTN_disp, xv_disp )
        end
    end
end


used_bits(x::Union{Signed, Unsigned}) = sizeof(x)*8 - leading_zeros(x)

"""
    kronecker_product(a, b)

Returns `a ⊗ b`, see: https://en.wikipedia.org/wiki/Kronecker_product.
"""
function kronecker_product(a::T, b::T) where T <: AbstractArray
    out = similar(a, size(a) .* size(b))
    for i in eachindex(IndexCartesian(), a)
        for j in eachindex(IndexCartesian(), b)
            out[CartesianIndex((Tuple(i) .- 1) .* size(b) .+ Tuple(j))] = a[i]*b[j]
        end
    end
    out
end

"""
    kronecker_power(a, power)

returns `a` kroneker producted with itself `power` times. Does not copy when power == 1.

See also: [`kronecker_product`](@ref)
"""
function kronecker_power(a::AbstractArray, power::Integer)
    if power < 0
        throw(ArgumentError("Power must be non-negative, got $power < 0."))
    elseif power == 0
        out = similar(a, ntuple(_->1, ndims(a)))
        out[firstindex(out)] = one(eltype(a))
        return out
    else power == 1
        out = a
        for i in 2:power
            out = kronecker_product(out, a)
        end
        out
    end
end


import Base: iterate

export qmr, qmr!

mutable struct LanczosDecomp{T, rT, matT, mat_tT, vecT}
    A::matT
    At::mat_tT

    v_prev::vecT # Orthonormal basis vectors for A
    v_curr::vecT # Orthonormal basis vectors for A
    v_next::vecT # Orthonormal basis vectors for A

    w_prev::vecT # Orthonormal basis vectors for A'
    w_curr::vecT # Orthonormal basis vectors for A'
    w_next::vecT # Orthonormal basis vectors for A'

    α::T
    β_prev::T
    β_curr::T
    δ::T
    resnorm::rT
end
function LanczosDecomp(x, A::matT, b; initially_zero = false) where {matT}
    T = eltype(A)

    # we choose:
    # v_0 = 0, v_1 = b - Ax
    # w_0 = 0, w_1 = v_1
    # v_2 and w_2 are uninitialized and used as workspace

    v_prev = zero(x)
    v_curr = copy(b)
    v_next = similar(x)

    if !initially_zero
        # r0 = A*x - b
        mul!(v_next, A, x)
        axpy!(-one(T), v_next, v_curr)
    end
    resnorm = norm(v_curr)
    rmul!(v_curr, inv(resnorm))

    w_prev = zero(x)
    w_curr = copy(v_curr)
    w_next = similar(x)

    α = zero(T)
    β_prev = zero(T)
    β_curr = zero(T)
    δ = zero(T)

    LanczosDecomp(
        A, adjoint(A),
        v_prev, v_curr, v_next,
        w_prev, w_curr, w_next,
        α, β_prev, β_curr, δ,
        resnorm)
end

start(::LanczosDecomp) = 1
done(l::LanczosDecomp, iteration::Int) = false
function iterate(l::LanczosDecomp, iteration::Int=start(l))
    # Following Algorithm 7.1 of Saad
    if done(l, iteration)
        return nothing
    end

    mul!(l.v_next, l.A, l.v_curr)

    l.α = dot(l.v_next, l.w_curr) # v_next currently contains A*v_curr
    axpy!(-conj(l.α), l.v_curr, l.v_next)
    if iteration > 1 # β_1 = 0
        axpy!(-conj(l.β_curr), l.v_prev, l.v_next)
    end

    mul!(l.w_next, l.At, l.w_curr)
    axpy!(-l.α, l.w_curr, l.w_next)
    if iteration > 1 # δ_1 = 0
        axpy!(-l.δ, l.w_prev, l.w_next)
    end

    vw = dot(l.v_next, l.w_next)
    l.δ = sqrt(abs(vw))
    if iszero(l.δ)
        return nothing
    end

    l.β_prev = l.β_curr
    l.β_curr = vw / l.δ

    rmul!(l.v_next, inv(l.δ))
    rmul!(l.w_next, inv(l.β_curr))

    l.w_next, l.w_curr, l.w_prev = l.w_prev, l.w_next, l.w_curr
    l.v_next, l.v_curr, l.v_prev = l.v_prev, l.v_next, l.v_curr

    return nothing, iteration + 1
end

mutable struct QMRIterable{T, xT, rT, lanczosT}
    x::xT

    lanczos::lanczosT
    resnorm::rT
    tol::rT
    maxiter::Int

    g::Vector{T} # right-hand size following Hessenberg multiplication
    H::Vector{T} # Hessenberg Matrix, computed on CPU

    c_prev::T
    c_curr::T
    s_prev::T
    s_curr::T

    p_prev::xT
    p_curr::xT
end

function qmr_iterable!(x, A, b;
                       abstol::Real = zero(real(eltype(b))),
                       reltol::Real = sqrt(eps(real(eltype(b)))),
                       maxiter::Int = size(A, 2),
                       initially_zero::Bool = false,
                       lookahead::Bool = false)
    T = eltype(x)

    lanczos = LanczosDecomp(x, A, b, initially_zero = initially_zero)

    resnorm = lanczos.resnorm
    g = [resnorm; zero(T)]
    H = zeros(T, 4)

    # Givens rotations
    c_prev, s_prev = one(T), zero(T)
    c_curr, s_curr = one(T), zero(T)

    # Projection operators
    p_prev = zero(x)
    p_curr = zero(x)

    tolerance = max(reltol * lanczos.resnorm, abstol)

    QMRIterable(x,
        lanczos, resnorm, tolerance,
        maxiter,
        g, H,
        c_prev, c_curr, s_prev, s_curr,
        p_prev, p_curr
    )
end

converged(q::QMRIterable) = q.resnorm ≤ q.tol
start(::QMRIterable) = 1
done(q::QMRIterable, iteration::Int) = iteration > q.maxiter || converged(q)

function iterate(q::QMRIterable, iteration::Int=start(q))
    # Check for termination first
    if done(q, iteration)
        return nothing
    end

    iterate(q.lanczos, iteration)

    # for classical Lanczos algorithm, only need 4 elements of last column
    # of Hessenberg matrix
    # H[1] is h_m,(m-2), H[2] is h_m,(m-1), H[3] is h_m,m, and H[4] is h_m,(m+1)
    q.H[2] = conj(q.lanczos.β_prev) # β_m
    q.H[3] = conj(q.lanczos.α)      # α_m
    q.H[4] = q.lanczos.δ            # δ_(m+1)

    # Rotation on H[1] and H[2]. Note that H[1] = 0 initially
    if iteration > 2
        q.H[1] = q.s_prev * q.H[2]
        q.H[2] = q.c_prev * q.H[2]
    end

    # Rotation on H[2] and H[3]
    if iteration > 1
        tmp    = -conj(q.s_curr) * q.H[2] + q.c_curr * q.H[3]
        q.H[2] =       q.c_curr  * q.H[2] + q.s_curr * q.H[3]
        q.H[3] = tmp
    end

    a, b = q.H[3], q.H[4]
    # Compute coefficients for Ω_m, and apply it
    c, s, q.H[3] = givensAlgorithm(q.H[3], q.H[4])

    # Apply Ω_m to rhs
    q.g[2] = -conj(s) * q.g[1]
    q.g[1] =       c  * q.g[1]

    # H is now properly factorized
    # compute p_m
    # we re-use q.lanczos.v_next as workspace for p_m
    q.lanczos.v_next .= q.lanczos.v_prev # we need v_m, not v_m+1
    iteration > 1 && axpy!(-q.H[2], q.p_curr, q.lanczos.v_next)
    iteration > 2 && axpy!(-q.H[1], q.p_prev, q.lanczos.v_next)
    rmul!(q.lanczos.v_next, inv(q.H[3]))

    # iterate x_n = x_n-1 + γ_m p_m
    axpy!(q.g[1], q.lanczos.v_next, q.x)

    # transfer coefficients of Givens rotations for next iteration
    q.c_prev, q.s_prev, q.c_curr, q.s_curr = q.c_curr, q.s_curr, c, s
    q.p_prev .= q.p_curr
    q.p_curr .= q.lanczos.v_next
    q.g[1] = q.g[2]

    # approximate residual
    # See Proposition 7.3 of Saad
    q.resnorm = abs(q.g[2])

    q.resnorm, iteration + 1
end

"""
    qmr(A, b; kwargs...) -> x, [history]

Same as [`qmr!`](@ref), but allocates a solution vector `x` initialized with zeros.
"""
qmr(A, b; kwargs...) = qmr!(zerox(A, b), A, b; initially_zero = true, kwargs...)

"""
    qmr!(x, A, b; kwargs...) -> x, [history]

Solves the problem ``Ax = b`` with the Quasi-Minimal Residual (QMR) method.

# Arguments
- `x`: Initial guess, will be updated in-place;
- `A`: linear operator;
- `b`: right-hand side.

## Keywords

- `initally_zero::Bool`: If `true` assumes that `iszero(x)` so that one
  matrix-vector product can be saved when computing the initial residual
  vector;
- `maxiter::Int = size(A, 2)`: maximum number of iterations;
- `abstol::Real = zero(real(eltype(b)))`,
  `reltol::Real = sqrt(eps(real(eltype(b))))`: absolute and relative
  tolerance for the stopping condition
  `|r_k| ≤ max(reltol * |r_0|, abstol)`, where `r_k = A * x_k - b`
- `log::Bool`: keep track of the residual norm in each iteration;
- `verbose::Bool`: print convergence information during the iteration.

# Return values

**if `log` is `false`**

- `x`: approximate solution.

**if `log` is `true`**

- `x`: approximate solution;

- `history`: convergence history.

[^Saad2003]:
    Saad, Y. (2003). Interactive method for sparse linear system.
[^Freund1990]:
    Freund, W. R., & Nachtigal, N. M. (1990). QMR : for a Quasi-Minimal
    Residual Linear Method Systems. (December).
"""
function qmr!(x, A, b;
              abstol::Real = zero(real(eltype(b))),
              reltol::Real = sqrt(eps(real(eltype(b)))),
              maxiter::Int = size(A, 2),
              lookahead::Bool = false,
              log::Bool = false,
              initially_zero::Bool = false,
              verbose::Bool = false)
    history = ConvergenceHistory(partial = !log)
    history[:abstol] = abstol
    history[:reltol] = reltol
    log && reserve!(history, :resnorm, maxiter)

    iterable = qmr_iterable!(x, A, b; abstol = abstol, reltol = reltol,
                             maxiter = maxiter, initially_zero = initially_zero)

    verbose && @printf("=== qmr ===\n%4s\t%7s\n","iter","resnorm")

    for (iteration, residual) = enumerate(iterable)
        if log
            nextiter!(history)
            push!(history, :resnorm, residual)
        end

        verbose && @printf("%3d\t%1.2e\n", iteration, residual)
    end

    verbose && println()
    log && setconv(history, converged(iterable))
    log && shrink!(history)

    log ? (x, history) : x
end


using Gadfly, Distributions

set_default_plot_size(6inch, 3inch)

vals = rand(MultivariateNormal([0.0, 0.0], [1.0 0.5; 0.5 1.0]), 10000)
plot(x=vals[1,:], y=vals[2,:], Geom.hexbin)


basis_blades(sig, grade=1) = [Blade{sig}(1, i) for i ∈ bits_of_grade(grade, dimension(sig))]

basis(sig; grade=1) = basis_blades(sig, grade)
basis(dim::Integer; grade=1) = basis(Tuple(fill(1, dim)); grade=1)
basis(osig::OffsetSignature{sig,indices}; grade=1) where {sig,indices} = OffsetArray(basis_blades(osig, grade), indices)

# return expression containing variable assignments for each basis blade for the given signature.
function generate_blades(combos, sig)
	bvs = collect(basis(sig)) # needed because `OffsetArrays` don't work with, e.g., `powerset`
	vars = Symbol[]
	vals = Blade{sig}[]
	for (ordered_bvs, ublade) ∈ zip(combos(bvs), combos(1:dimension(sig)))
		varname = basis_blade_label(sig, ublade)
		isempty(varname) && continue
		push!(vars, Symbol(varname))
		push!(vals, prod(ordered_bvs))
	end
	quote
		@info "Defined basis blades $(join($vars, ", "))"
		$([ :($(esc(k)) = $(v)) for (k, v) ∈ zip(vars, vals) ]...)
		nothing
	end
end
generate_blades(sig) = generate_blades(x -> [[i] for i ∈ x], sig)

parse_sig(n::Integer) = Tuple(ones(n))
parse_sig(sig::MetricSignature) = sig
parse_sig(sig::Tuple{Vararg{<:Integer}}) = sig
parse_sig(sig::NamedTuple) = sig
parse_sig(labels::Tuple{Vararg{Symbol}}) = NamedTuple{labels}(ones(Int, length(labels)))
parse_sig(sig::AbstractString) = Tuple(get(Dict('+'=>1, '-'=>-1, '0'=>0), i) do i
	throw(ArgumentError("Invalid character in signature string: $i"))
end for i ∈ sig)
parse_sig(sig) = throw(ArgumentError("could not interpret $sig as metric signature"))

function parse_sig(sig::Tuple)
	# signature specified by series of symbols/exprs `label[=square]`
	labels = Symbol[]
	squares = []
	for bv ∈ sig
		if bv isa Symbol
			label, sq = bv, 1
		elseif bv isa Expr && bv.head == :(=)
			label, sq = bv.args
		else
			throw(ArgumentError("Invalid term '$bv' in signature specification"))
		end
		push!(labels, label)
		push!(squares, sq)
	end
	NamedTuple{tuple(labels...)}(squares)
end

# for a single argument `@basis X`, interpret `X` as an expression which evaluates to a signature.
# for many arguments, `@basis x1 x2 x3...`, interpret `(x1, x2, x3)` as a signature.
parse_sig(::Module, args...) = parse_sig(args)
parse_sig(m::Module, args) = let a = @eval m $args
	parse_sig(a)
end

"""
	@basis sig
	@basis x y z...

Populate the local namespace with basis blades of all grades for
the geometric algebra with signature `sig`, which may be provided
as a string, tuple, named tuple or `<:AbstractMetricSignature`.
If multiple arguments are given, each is interpreted as specifying
the label and square (defaulting to `1`) of a basis vector.
For an `n`-dimensional signature, `2^n` basis blades are defined.

See also `@basisall`.

Examples
===
```jldoctest
julia> @basis "+++"
[ Info: Defined basis blades v, v1, v2, v3, v12, v13, v23, v123

julia> @basis (+1,0,-1)
[ Info: Defined basis blades v, v1, v2, v3, v12, v13, v23, v123

julia> @basis x y
[ Info: Defined basis blades x, y, xy

julia> @basis dt=-1 dx=1 dy
[ Info: Defined basis blades dt, dx, dy, dtdx, dtdy, dxdy, dtdxdy

julia> dt^2 + dxdy*dy
MixedMultivector{⟨dt-,dx+,dy+⟩, Vector{Int64}}:
 -1
 1 dx
```
"""
macro basis(a, args...)
	generate_blades(powerset, parse_sig(__module__, a, args...))
end

"""
	@basisall

Similarly to `@basis`, populate the local namespace with basis blades,
but including all permutations of the basis blades.

Note that more than `2^n` variables are defined for a signature of dimension `n`!

Example
===
```jldoctest
julia> @basisall x y z
[ Info: Defined basis blades x, y, z, xy, yx, xz, zx, yz, zy, xyz, xzy, yxz, yzx, zxy, zyx

julia> zyx
Grade-3 Blade{⟨x+,y+,z+⟩, 3, 0b111, Int64}:
 -1 xyz
```
"""
macro basisall(a, args...)
	generate_blades(parse_sig(__module__, a, args...)) do bvs
		Iterators.flatten(permutations.(powerset(bvs)))
	end
end


using Test
using LowRankModels

## Tests for NNDSVD initialization

# apply init_nnmf! to random dataset
m,n,k = 10,5,3
A = rand(m,k)*rand(k,n)
losses = fill(QuadLoss(),n)
r = NonNegConstraint()
glrm = GLRM(A,losses,r,r,k)
init_nndsvd!(glrm)

# Test dims and nonnegativity of X,Y
@test(all(glrm.X .>= 0.0))
@test(all(glrm.Y .>= 0.0))
@test(size(glrm.X,1) == k)
@test(size(glrm.X,2) == m)
@test(size(glrm.Y,1) == k)
@test(size(glrm.Y,2) == n)

# validate against NMF.jl
A = [ 0.245774   0.481246   0.293614   0.528272   0.608255
      0.906146   0.0847498  0.963121   0.113728   0.552843
      0.269553   0.830981   0.067466   0.854045   0.170701
      0.781552   0.0302068  0.709045   0.578069   0.542792
      0.824889   0.538719   0.881363   0.229199   0.356273
      0.0530284  0.618962   0.96237    0.0877032  0.921746
      0.295122   0.626784   0.348475   0.299937   0.35043
      0.0499904  0.728344   0.573141   0.850758   0.425369
      0.872088   0.322181   0.903238   0.695946   0.706841
      0.542786   0.581426   0.0477561  0.601374   0.176598 ]

# Wt and H produced by NMF.jl
Wt = [0.25814 0.34329 0.24898 0.33489 0.35953 0.33798 0.23164 0.31705 0.43949 0.22484
      0.24305 0.0 0.71244 0.0 0.0 0.0 0.19987 0.49127 0.0 0.38994
      0.13531 0.0 0.0 0.0 0.0 0.94666 0.13278 0.26056 0.0 0.0]

H = [1.60738 1.42165 1.97295 1.47396 1.57709
     0.0 0.67822 0.0 0.6408 0.0
     0.0 0.23296 0.27065 0.0 0.35102]

# initialize glrm and check output
glrm.A = A
init_nndsvd!(glrm; scale=false)

@test(all(round(glrm.X,5) .== Wt))
@test(all(round(glrm.Y,5) .== H))

# Test with missing entries for A
obs = [(1,1),(1,3),(1,5),
       (2,1),(2,2),(2,4),
       (3,2),(3,3),(3,5),
       (4,1),(4,2),(4,5),
       (5,2),(5,3),(5,5),
       (6,1),(6,2),(6,5),
       (7,1),(7,3),(7,4),
       (8,2),(8,3),(8,4),
       (9,1),(9,3),(9,5),
       (10,3),(10,4),(10,5)]
glrm = GLRM(A,losses,r,r,k,obs=obs)
init_nndsvd!(glrm; max_iters=5)