# This file is part of TaylorSeries.jl, MIT licensed
#

using TaylorSeries

using Test
using LinearAlgebra

@testset "Tests for HomogeneousPolynomial and TaylorN" begin
    eeuler = Base.MathConstants.e

    @test HomogeneousPolynomial <: AbstractSeries
    @test HomogeneousPolynomial{Int} <: AbstractSeries{Int}
    @test TaylorN{Float64} <: AbstractSeries{Float64}

    set_variables("x", numvars=2, order=6)
    _taylorNparams = TaylorSeries.ParamsTaylorN(6, 2, String["x₁", "x₂"])
    @test _taylorNparams.order == get_order()
    @test _taylorNparams.num_vars == get_numvars()
    @test _taylorNparams.variable_names == get_variable_names()
    @test _taylorNparams.variable_symbols == get_variable_symbols()

    @test eltype(set_variables(Int, "x", numvars=2, order=6)) == TaylorN{Int}
    @test eltype(set_variables("x", numvars=2, order=6)) == TaylorN{Float64}
    @test eltype(set_variables(BigInt, "x y", order=6)) == TaylorN{BigInt}
    @test eltype(set_variables("x y", order=6)) == TaylorN{Float64}
    @test eltype(set_variables(Int, :x, numvars=2, order=6)) == TaylorN{Int}
    @test eltype(set_variables(:x, numvars=2, order=6)) == TaylorN{Float64}
    @test eltype(set_variables(BigInt, [:x,:y], order=6)) == TaylorN{BigInt}
    @test eltype(set_variables([:x,:y], order=6)) == TaylorN{Float64}
    @test typeof(show_params_TaylorN()) == Nothing
    @test typeof(show_monomials(2)) == Nothing

    @test TaylorSeries.coeff_table[2][1] == [1,0]
    @test TaylorSeries.index_table[2][1] == 7
    @test TaylorSeries.in_base(get_order(),[2,1]) == 15
    @test TaylorSeries.pos_table[4][15] == 2

    @test get_order() == 6
    @test get_numvars() == 2

    @test get_variables()[1].order == get_order()
    @test get_variables(2)[1].order == 2
    @test get_variables(3)[1] == TaylorN(1, order=3)
    @test get_variables(Int, 3)[1] == TaylorN(Int, 1, order=3)
    @test length(get_variables()) == get_numvars()

    x, y = set_variables("x y", order=6)
    @test axes(x) == axes(y) == ()
    @test axes(x[1]) == axes(y[2]) == ()
    @test size(x) == (7,)
    @test size(x[1]) == (2,)
    @test size(x[2]) == (3,)
    @test firstindex(x) == 0
    @test firstindex(x[end]) == 1
    @test lastindex(y) == get_order()
    @test eachindex(x) == 0:6
    @test iterate(x) == (HomogeneousPolynomial([0.0], 0), 1)
    @test iterate(y, 1) == (HomogeneousPolynomial([0.0, 1.0], 1), 2)
    @test iterate(x, 7) == nothing

    @test x.order == 6
    @test TaylorSeries.name_taylorNvar(1) == " x"
    @test TaylorSeries._params_TaylorN_.variable_names == ["x","y"]
    @test TaylorSeries._params_TaylorN_.variable_symbols == [:x, :y]
    @test get_variable_symbols() == [:x, :y]
    @test TaylorSeries.lookupvar(:x) == 1
    @test TaylorSeries.lookupvar(:α) == 0
    @test TaylorSeries.get_variable_names() == ["x", "y"]
    @test x == HomogeneousPolynomial(Float64, 1)
    @test x == HomogeneousPolynomial(1)
    @test y == HomogeneousPolynomial(Float64, 2)
    @test y == HomogeneousPolynomial(2)
    @test !isnan(x)

    set_variables("x", numvars=2, order=17)
    v = [1,2]
    @test typeof(TaylorSeries.resize_coeffsHP!(v,2)) == Nothing
    @test v == [1,2,0]
    @test_throws AssertionError TaylorSeries.resize_coeffsHP!(v,1)
    hpol_v = HomogeneousPolynomial(v)
    hpol_v[3] = 3
    @test v == [1,2,3]
    hpol_v[1:3] = 3
    @test v == [3,3,3]
    hpol_v[1:2:2] = 0
    @test v == [0,3,3]
    hpol_v[1:1:2] = [1,2]
    @test all(hpol_v[1:1:2] .== [1,2])
    @test v == [1,2,3]
    hpol_v[:] = zeros(Int, 3)
    @test hpol_v == 0

    tn_v = TaylorN(HomogeneousPolynomial(zeros(Int, 3)))
    tn_v[0] = 1
    @test tn_v == 1
    tn_v[0:1] = [0, 1]
    @test tn_v[0] == 0 && tn_v[1] == HomogeneousPolynomial(1, 1)
    tn_v[0:1] = [HomogeneousPolynomial(0, 0), HomogeneousPolynomial([0,1])]
    @test tn_v[0] == 0 && tn_v[1] == HomogeneousPolynomial([0,1], 1)
    tn_v[:] = [HomogeneousPolynomial(1, 0), HomogeneousPolynomial(0, 1), hpol_v]
    @test tn_v == 1
    tn_v[:] = 0
    @test tn_v == 0
    tn_v[:] = [3,1,0]
    @test tn_v == TaylorN([HomogeneousPolynomial(3, 0), HomogeneousPolynomial(1, 1)], 2)
    tn_v[0:2] = [HomogeneousPolynomial(3, 0), HomogeneousPolynomial(1, 1), HomogeneousPolynomial(0, 2)]
    @test tn_v == TaylorN([HomogeneousPolynomial(3, 0), HomogeneousPolynomial(1, 1)], 2)
    tn_v[0:2:2] = [0,0]
    @test tn_v == TaylorN(HomogeneousPolynomial(1, 1), 2)

    xH = HomogeneousPolynomial([1,0])
    yH = HomogeneousPolynomial([0,1],1)
    @test xH == convert(HomogeneousPolynomial{Float64},xH)
    @test HomogeneousPolynomial(0,0)  == 0
    xT = TaylorN(xH, 17)
    yT = TaylorN(Int, 2, order=17)
    zeroT = zero( TaylorN([xH],1) )
    @test zeroT.coeffs == zeros(HomogeneousPolynomial{Int}, 1)
    @test size(xH) == (2,)
    @test firstindex(xH) == 1
    @test lastindex(yH) == 2
    @test length(zeros(HomogeneousPolynomial{Int}, 1)) == 2
    @test one(HomogeneousPolynomial(1,1)) == HomogeneousPolynomial([1,1])
    uT = one(convert(TaylorN{Float64},yT))
    @test uT == one(HomogeneousPolynomial)
    @test uT == convert(TaylorN{Float64},uT)
    @test zeroT[0] == HomogeneousPolynomial(0, 0)
    @test uT[0] == HomogeneousPolynomial(1, 0)
    @test ones(xH,1) == [1, xH+yH]
    @test typeof(ones(xH,2)) == Array{HomogeneousPolynomial{Int},1}
    @test length(ones(xH,2)) == 3
    @test ones(HomogeneousPolynomial{Complex{Int}},0) ==
        [HomogeneousPolynomial([complex(1,0)], 0)]
    @test !isnan(uT)
    @test TaylorSeries.fixorder(xH,yH) == (xH,yH)
    @test_throws AssertionError TaylorSeries.fixorder(zeros(xH,0)[1],yH)

    @test constant_term(xT) == 0
    @test constant_term(uT) == 1.0
    @test constant_term(xT) == constant_term(yT)
    @test constant_term(xH) == xH
    @test linear_polynomial(xT) == xT
    @test linear_polynomial(xT+yT) == xT+yT
    @test linear_polynomial(uT) == zero(yT)

    @test get_order(zeroT) == 1
    @test xT[1][1] == 1
    @test yH[2] == 1
    @test getcoeff(xT,(1,0)) == getcoeff(xT,[1,0]) == 1
    @test getcoeff(yH,(1,0)) == getcoeff(yH,[1,0]) == 0
    @test typeof(convert(HomogeneousPolynomial,1im)) ==
        HomogeneousPolynomial{Complex{Int}}
    @test convert(HomogeneousPolynomial,1im) ==
        HomogeneousPolynomial([complex(0,1)], 0)
    @test convert(HomogeneousPolynomial{Int},[1,1]) == xH+yH
    @test convert(HomogeneousPolynomial{Float64},[2,-1]) == 2.0xH-yH
    @test typeof(convert(TaylorN,1im)) == TaylorN{Complex{Int}}
    @test convert(TaylorN, 1im) ==
        TaylorN([HomogeneousPolynomial([complex(0,1)], 0)], 0)
    @test convert(TaylorN{Float64}, yH) == 1.0*yT
    @test convert(TaylorN{Float64}, [xH,yH]) == xT+1.0*yT
    @test convert(TaylorN{Int}, [xH,yH]) == xT+yT
    @test promote(xH, [1,1])[2] == xH+yH
    @test promote(xH, yT)[1] == xT
    @test promote(xT, [xH,yH])[2] == xT+yT
    @test typeof(promote(im*xT,[xH,yH])[2]) == TaylorN{Complex{Int}}

    @test iszero(zeroT.coeffs)
    @test iszero(zero(xH))
    @test !iszero(uT)
    @test iszero(zeroT)

    @test eltype(xH) == Int
    @test length(xH) == 2
    @test zero(xH) == 0*xH
    @test one(yH) == xH+yH
    @test xH * true == xH
    @test false * yH == zero(yH)
    @test get_order(yH) == 1
    @test get_order(xT) == 17
    @test xT * true == xT
    @test false * yT == zero(yT)
    @test HomogeneousPolynomial([1.0])*xH == xH

    @test xT == TaylorN([xH])
    @test one(xT) == TaylorN(1,5)
    @test TaylorN(uT) == convert(TaylorN{Complex},1)
    @test get_numvars() == 2
    @test length(uT) == get_order()+1
    @test eltype(convert(TaylorN{Complex{Float64}},1)) == Complex{Float64}

    @test 1+xT+yT == TaylorN(1,1) + TaylorN([xH,yH],1)
    @test xT-yT-1 == TaylorN([-1,xH-yH])
    @test xT*yT == TaylorN([HomogeneousPolynomial([0,1,0],2)])
    @test (1/(1-xT))[3] == HomogeneousPolynomial([1.0],3)
    @test xH^20 == HomogeneousPolynomial([0], get_order())
    @test (yT/(1-xT))[4] == xH^3 * yH
    @test mod(1+xT,1) == +xT
    @test (rem(1+xT,1))[0] == 0
    @test mod(1+xT,1.0) == +xT
    @test (rem(1+xT,1.0))[0] == 0
    @test abs(1-xT)  == 1-xT
    @test abs(-1-xT)  == 1+xT
    @test differentiate(yH,1) == differentiate(xH, :x₂)
    @test derivative(mod2pi(2pi+yT^3),2) == derivative(yT^3, :x₂)
    @test derivative(yT^3, :x₂) == derivative(yT^3, (0,1))
    @test derivative(yT) == zeroT == derivative(yT, (1,0))
    @test derivative((0,1), yT) == 1
    @test -xT/3im == im*xT/3
    @test (xH/3im)' == im*xH/3
    @test xT/BigInt(3) == TaylorN(BigFloat,1)/3
    @test xT/complex(0,BigInt(3)) == -im*xT/BigInt(3)
    @test (xH/complex(0,BigInt(3)))' ==
        im*HomogeneousPolynomial([BigInt(1),0])/3
    @test evaluate(xH) == zero(eltype(xH))
    @test xH() == zero(eltype(xH))
    @test xH([1,1]) == evaluate(xH, [1,1])
    @test xH((1,1)) == xH(1, 1.0) == evaluate(xH, (1, 1.0)) == 1
    hp = -5.4xH+6.89yH
    @test hp([1,1]) == evaluate(hp, [1,1])
    vr = rand(2)
    @test hp(vr) == evaluate(hp, vr)

    @test integrate(yH,1) == integrate(xH, :x₂)
    p = (xT-yT)^6
    @test integrate(derivative(p, 1), 1, yT^6) == p
    @test integrate(derivative(p, :x₁), :x₁, yT^6) == p
    @test derivative(integrate(p, 2), 2) == p
    @test derivative(integrate(p, :x₂), :x₂) == p
    @test derivative(TaylorN(1.0, get_order())) == TaylorN(0.0, get_order())
    @test integrate(TaylorN(6.0, get_order()), 1) == 6xT
    @test integrate(TaylorN(0.0, get_order()), 2) == TaylorN(0.0, get_order())
    @test integrate(TaylorN(0.0, get_order()), 2, xT) == xT
    @test integrate(TaylorN(0.0, get_order()), :x₂, xT) == xT
    @test integrate(xT^17, 2) == TaylorN(0.0, get_order())
    @test integrate(xT^17, 1, yT) == yT
    @test integrate(xT^17, 1, 2.0) == TaylorN(2.0, get_order())
    @test integrate(xT^17, :x₁, 2.0) == TaylorN(2.0, get_order())
    @test_throws AssertionError integrate(xT, 1, xT)
    @test_throws AssertionError integrate(xT, :x₁, xT)
    @test_throws AssertionError derivative(xT, (1,))
    @test_throws AssertionError derivative(xT, (1,2,3))
    @test_throws AssertionError derivative(xT, (-1,2))
    @test_throws AssertionError derivative((1,), xT)
    @test_throws AssertionError derivative((1,2,3), xT)
    @test_throws AssertionError derivative((-1,2), xT)

    @test derivative(2xT*yT^2, (8,8)) == 0
    @test derivative((8,8), 2xT*yT^2) == 0
    @test derivative(2xT*yT^2, 1) == 2yT^2
    @test derivative((1,0), 2xT*yT^2) == 0
    @test derivative(2xT*yT^2, (1,2)) == 4*one(yT)
    @test derivative((1,2), 2xT*yT^2) == 4
    @test xT*xT^3 == xT^4
    txy = 1.0 + xT*yT - 0.5*xT^2*yT + (1/3)*xT^3*yT + 0.5*xT^2*yT^2
    @test getindex((1+TaylorN(1))^TaylorN(2),0:4) == txy.coeffs[1:5]
    @test ( (1+TaylorN(1))^TaylorN(2) )[:] == ( (1+TaylorN(1))^TaylorN(2) ).coeffs[:]
    @test txy.coeffs[:] == txy[:]
    @test txy.coeffs[:] == txy[0:end]
    txy[:] .= ( -1.0 + 3xT*yT - xT^2*yT + (4/3)*xT^3*yT + (1/3)*xT*yT^3 + 0.5*xT^2*yT^2 + 0.5*xT*yT^3 )[:]
    @test txy[:] == ( -1.0 + 3xT*yT - xT^2*yT + (4/3)*xT^3*yT + (1/3)*xT*yT^3 + 0.5*xT^2*yT^2 + 0.5*xT*yT^3 )[:]
    txy[2:end-1] .= ( 1.0 - xT*yT + 0.5*xT^2*yT - (2/3)*xT*yT^3 - 0.5*xT^2*yT^2  + 7*xT^3*yT )[2:end-1]
    @test txy[2:end-1] == ( 1.0 - xT*yT + 0.5*xT^2*yT - (2/3)*xT*yT^3 - 0.5*xT^2*yT^2  + 7*xT^3*yT )[2:end-1]

    a = -5.0 + sin(xT+yT^2)
    b = deepcopy(a)
    @test a[:] == a[0:end]
    @test a[:] == b[:]
    @test a[1:end] == b[1:end]
    @test a[end][:] == b[end][:]
    @test a[end][1:end] == b[end][1:end]
    a[end][:] .= rand.()
    rv = a[end][:]
    @test a[end][:] == rv
    @test a[end][:] != b[end][:]
    a[end][1:end] .= rand.()
    rv = a[end][1:end]
    @test a[end][1:end] == rv
    @test a[end][1:end] != b[end][1:end]
    @test a[0:2:end] == a.coeffs[1:2:end]
    a[0:1:end] .= 0.0
    @test a == zero(a)

    hp = HomogeneousPolynomial(1)^8
    rv1 = rand( length(hp) )
    hp[:] = rv1
    @test rv1 == hp[:]
    rv2 = rand( length(hp)-2 )
    hp[1:end-2] .= rv2
    @test hp[1:end-2] == rv2
    @test hp[end-1:end] == rv1[end-1:end]
    hp[3:4] .= 0.0
    @test hp[1:2] == rv2[1:2]
    @test hp[3:4] == zeros(2)
    @test hp[5:end-2] == rv2[5:end]
    @test hp[end-1:end] == rv1[end-1:end]
    hp[:] = 0.0
    @test hp[:] == zero(rv1)
    @test all(hp[end-1:1:end] .== 0.0)

    pol = sin(xT+yT*xT)+yT^2-(1-xT)^3
    q = deepcopy(pol)
    q[:] = 0.0
    @test get_order.(q[:]) == collect(0:q.order)
    @test q[:] == zero(q[:])
    q[:] .= pol.coeffs
    @test q == pol
    @test q[:] == pol[:]
    q[2:end-1] .= 0.0
    @test q[2:end-1] == zero.(q[2:end-1])
    @test q[1] == pol[1]
    @test q[end] == pol[end]
    # q[:] = pol.coeffs
    # zH0 = zero(HomogeneousPolynomial{Float64})
    q[:] = 1.0
    @test q[1] == HomogeneousPolynomial([1,0])
    @test q[2] == HomogeneousPolynomial([1,0,0])
    q[:] .= pol.coeffs
    q[2:end-1] = one.(q[2:end-1])
    @test q[2:end-1] == one.(q[2:end-1])
    @test q[2] == HomogeneousPolynomial([1,1,1])
    @test q[1] == pol[1]
    @test q[end] == pol[end]
    q[:] .= pol.coeffs
    zHall = zeros(HomogeneousPolynomial{Float64}, q.order)
    q[:] .= zHall
    @test q[:] == zHall
    q[:] .= pol.coeffs
    q[1:end-1] .= zHall[2:end-1]
    @test q[1:end-1] == zHall[2:end-1]
    q[:] .= pol.coeffs
    @test q[:] != zeros(q.order+1)
    q[:] .= zeros(q.order+1)
    @test q[:] == zeros(q.order+1)
    q[:] .= pol.coeffs
    q[1:end-1] .= zeros(q.order+1)[2:end-1]
    @test q != pol
    @test all(q[1:1:end-1] .== 0.0)
    @test q[1:end-1] == zeros(q.order+1)[2:end-1]
    @test q[0] == pol[0]
    @test q[end] == pol[end]
    q[:] .= pol.coeffs
    pol2 = cos(sin(xT)-yT^3*xT)-3yT^2+sqrt(1-xT)
    q[2:end-2] .= pol2.coeffs[3:end-2]
    @test q[0:1] == pol[0:1]
    @test q[2:end-2] == pol2[2:end-2]
    @test q[end-1:end] == pol[end-1:end]
    @test q[2:2:end-2] == pol2[2:2:end-2]
    @test q[end-1:1:end] == pol[end-1:1:end]
    q[end-2:2:end] .= [0.0, 0.0]
    @test q[end-2] == 0.0
    @test_throws AssertionError q[end-2:2:end] = [0.0, 0.0, 0.0]
    q[end-2:2:end] .= pol.coeffs[end-2:2:end]
    @test q[end-2] == pol[end-2]
    q[end-2:2:end] .= pol.coeffs[end-2:2:end]
    @test_throws AssertionError q[end-2:2:end] = pol.coeffs[end-1:2:end]

    @test_throws AssertionError yT^(-2)
    @test_throws AssertionError yT^(-2.0)
    @test (1+xT)^(3//2) == ((1+xT)^0.5)^3
    @test real(xH) == xH
    @test imag(xH) == zero(xH)
    @test conj(im*yH) == (im*yH)'
    @test conj(im*yT) == (im*yT)'
    @test real( exp(1im * xT)) == cos(xT)
    @test getcoeff(convert(TaylorN{Rational{Int}},cos(xT)),(4,0)) ==
        1//factorial(4)
    cr = convert(TaylorN{Rational{Int}},cos(xT))
    @test getcoeff(cr,(4,0)) == 1//factorial(4)
    @test imag((exp(yT))^(-1im)') == sin(yT)
    exy = exp( xT+yT )
    @test evaluate(exy) == 1
    @test evaluate(exy, 0.1im, 0.01im) == exp(0.11im)
    @test exy(0.1im, 0.01im) == exp(0.11im)
    @test evaluate(exy,(0.1im, 0.01im)) == exp(0.11im)
    @test exy((0.1im, 0.01im)) == exp(0.11im)
    @test exy(true, (0.1im, 0.01im)) == exp(0.11im)
    @test evaluate(exy, (0.1im, 0.01im), sorting=false) == exy(false, (0.1im, 0.01im))
    @test evaluate(exy, (0.1im, 0.01im), sorting=false) == exy(false, 0.1im, 0.01im)
    @test evaluate(exy,[0.1im, 0.01im]) == exp(0.11im)
    @test exy([0.1im, 0.01im]) == exp(0.11im)
    @test isapprox(evaluate(exy, (1,1)), eeuler^2)
    @test exy(:x₁, 0.0) == exp(yT)
    txy = tan(xT+yT)
    @test getcoeff(txy,(8,7)) == 929569/99225
    ptxy = xT + yT + (1/3)*( xT^3 + yT^3 ) + xT^2*yT + xT*yT^2
    @test getindex(tan(TaylorN(1)+TaylorN(2)),0:4) == ptxy.coeffs[1:5]
    @test evaluate(xH*yH, 1.0, 2.0) == (xH*yH)(1.0, 2.0) == 2.0
    @test evaluate(xH*yH, (1.0, 2.0)) == 2.0
    @test evaluate(xH*yH, [1.0, 2.0]) == 2.0
    @test ptxy(:x₁, -1.0) == -1 + yT + (-1.0+yT^3)/3 + yT - yT^2
    @test ptxy(:x₁ => -1.0) == -1 + yT + (-1.0+yT^3)/3 + yT - yT^2
    @test evaluate(ptxy, :x₁ => -1.0) == -1 + yT + (-1.0+yT^3)/3 + yT - yT^2
    @test evaluate(ptxy, :x₁, -1.0) == -1 + yT + (-1.0+yT^3)/3 + yT - yT^2
    v = zeros(Int, 2)
    @test evaluate!([xT, yT], ones(Int, 2), v) == nothing
    @test v == ones(2)
    @test evaluate!([xT, yT][1:2], ones(Int, 2), v) == nothing
    @test v == ones(2)
    A_TN = [xT 2xT 3xT; yT 2yT 3yT]
    @test evaluate(A_TN, ones(2)) == [1.0 2.0 3.0; 1.0 2.0 3.0]
    @test evaluate(A_TN) == [0.0 0.0 0.0; 0.0 0.0 0.0]
    @test A_TN() == [0.0  0.0  0.0; 0.0  0.0  0.0]
    @test (view(A_TN,:,:))() == [0.0 0.0 0.0; 0.0 0.0 0.0]
    t = Taylor1(10)
    @test A_TN([t,t^2]) == [t 2t 3t; t^2 2t^2 3t^2]
    @test view(A_TN, :, :)(ones(2)) == A_TN(ones(2))
    @test view(A_TN, :, 1)(ones(2)) == A_TN[:,1](ones(2))

    @test evaluate(sin(asin(xT+yT)), [1.0,0.5]) == 1.5
    @test evaluate(asin(sin(xT+yT)), [1.0,0.5]) == 1.5
    @test tan(atan(xT+yT)) == xT+yT
    @test atan(tan(xT+yT)) == xT+yT
    @test atan(sin(1+xT+yT), cos(1+xT+yT)) == atan(sin(1+xT+yT)/cos(1+xT+yT))
    @test constant_term(atan(sin(3pi/4+xT+yT), cos(3pi/4+xT+yT))) == 3pi/4
    @test asin(xT+yT) + acos(xT+yT) == pi/2

    @test -sinh(xT+yT) + cosh(xT+yT) == exp(-(xT+yT))
    @test  sinh(xT+yT) + cosh(xT+yT) == exp(xT+yT)
    @test evaluate(- sinh(xT+yT)^2 + cosh(xT+yT)^2 , rand(2)) == 1
    @test evaluate(- sinh(xT+yT)^2 + cosh(xT+yT)^2 , zeros(2)) == 1
    @test tanh(xT + yT + 0im) == -1im * tan((xT+yT)*1im)
    @test cosh(xT+yT) == real(cos(im*(xT+yT)))
    @test sinh(xT+yT) == imag(sin(im*(xT+yT)))

    xx = 1.0*zeroT
    TaylorSeries.add!(xx, 1.0*xT, 2yT, 1)
    @test xx[1] == HomogeneousPolynomial([1,2])
    TaylorSeries.add!(xx, 5.0, 0)
    @test xx[0] == HomogeneousPolynomial([5.0])
    TaylorSeries.add!(xx, -5.0, 1)
    @test xx[1] == zero(xx[1])
    TaylorSeries.subst!(xx, 1.0*xT, yT, 1)
    @test xx[1] == HomogeneousPolynomial([1,-1])
    TaylorSeries.subst!(xx, 5.0, 0)
    @test xx[0] == HomogeneousPolynomial([-5.0])
    TaylorSeries.subst!(xx, -5.0, 1)
    @test xx[1] == zero(xx[end])
    TaylorSeries.div!(xx, 1.0+xT, 1.0+xT, 0)
    @test xx[0] == HomogeneousPolynomial([1.0])
    TaylorSeries.pow!(xx, 1.0+xT, 0.5, 1)
    @test xx[1] == HomogeneousPolynomial([0.5,0.0])
    xx = 1.0*zeroT
    TaylorSeries.pow!(xx, 1.0+xT, 1.5, 0)
    @test xx[0] == HomogeneousPolynomial([1.0])
    TaylorSeries.pow!(xx, 1.0+xT, 1.5, 1)
    @test xx[1] == HomogeneousPolynomial([1.5,0.0])
    xx = 1.0*zeroT
    TaylorSeries.pow!(xx, 1.0+xT, 0, 0)
    @test xx[0] == HomogeneousPolynomial([1.0])
    TaylorSeries.pow!(xx, 1.0+xT, 1, 1)
    @test xx[1] == HomogeneousPolynomial([1.0,0.0])
    xx = 1.0*zeroT
    TaylorSeries.pow!(xx, 1.0+xT, 2, 0)
    @test xx[0] == HomogeneousPolynomial([1.0])
    TaylorSeries.pow!(xx, 1.0+xT, 2, 1)
    @test xx[1] == HomogeneousPolynomial([2.0,0.0])
    xx = 1.0*zeroT
    TaylorSeries.sqrt!(xx, 1.0+xT, 0)
    TaylorSeries.sqrt!(xx, 1.0+xT, 1)
    @test xx[0] == 1.0
    @test xx[1] == HomogeneousPolynomial([0.5,0.0])
    xx = 1.0*zeroT
    TaylorSeries.exp!(xx, 1.0*xT, 0)
    TaylorSeries.exp!(xx, 1.0*xT, 1)
    @test xx[0] == 1.0
    @test xx[1] == HomogeneousPolynomial([1.0,0.0])
    xx = 1.0*zeroT
    TaylorSeries.log!(xx, 1.0+xT, 0)
    TaylorSeries.log!(xx, 1.0+xT, 1)
    @test xx[0] == 0.0
    @test xx[1] == HomogeneousPolynomial(1.0,1)
    xx = 1.0*zeroT
    cxx = zero(xx)
    TaylorSeries.sincos!(xx, cxx, 1.0*xT, 0)
    TaylorSeries.sincos!(xx, cxx, 1.0*xT, 1)
    @test xx[0] == 0.0
    @test xx[1] == HomogeneousPolynomial(1.0,1)
    @test cxx[0] == 1.0
    @test cxx[1] == HomogeneousPolynomial(0.0,1)
    xx = 1.0*zeroT
    cxx = zero(xx)
    TaylorSeries.tan!(xx, 1.0*xT, cxx, 0)
    TaylorSeries.tan!(xx, 1.0*xT, cxx, 1)
    @test xx[0] == 0.0
    @test xx[1] == HomogeneousPolynomial(1.0,1)
    @test cxx[0] == 0.0
    @test cxx[1] == HomogeneousPolynomial(0.0,1)
    xx = 1.0*zeroT
    cxx = zero(xx)
    TaylorSeries.asin!(xx, 1.0*xT, cxx, 0)
    TaylorSeries.asin!(xx, 1.0*xT, cxx, 1)
    @test xx[0] == 0.0
    @test xx[1] == HomogeneousPolynomial(1.0,1)
    @test cxx[0] == 1.0
    @test cxx[1] == HomogeneousPolynomial(0.0,1)
    xx = 1.0*zeroT
    cxx = zero(xx)
    TaylorSeries.acos!(xx, 1.0*xT, cxx, 0)
    TaylorSeries.acos!(xx, 1.0*xT, cxx, 1)
    @test xx[0] == acos(0.0)
    @test xx[1] == HomogeneousPolynomial(-1.0,1)
    @test cxx[0] == 1.0
    @test cxx[1] == HomogeneousPolynomial(0.0,1)
    xx = 1.0*zeroT
    cxx = zero(xx)
    TaylorSeries.atan!(xx, 1.0*xT, cxx, 0)
    TaylorSeries.atan!(xx, 1.0*xT, cxx, 1)
    @test xx[0] == 0.0
    @test xx[1] == HomogeneousPolynomial(1.0,1)
    @test cxx[0] == 1.0
    @test cxx[1] == HomogeneousPolynomial(0.0,1)
    xx = 1.0*zeroT
    cxx = zero(xx)
    TaylorSeries.sinhcosh!(xx, cxx, 1.0*xT, 0)
    TaylorSeries.sinhcosh!(xx, cxx, 1.0*xT, 1)
    @test xx[0] == 0.0
    @test xx[1] == HomogeneousPolynomial(1.0,1)
    @test cxx[0] == 1.0
    @test cxx[1] == HomogeneousPolynomial(0.0,1)
    xx = 1.0*zeroT
    cxx = zero(xx)
    TaylorSeries.tanh!(xx, 1.0*xT, cxx, 0)
    TaylorSeries.tanh!(xx, 1.0*xT, cxx, 1)
    @test xx[0] == 0.0
    @test xx[1] == HomogeneousPolynomial(1.0,1)
    @test cxx[0] == 0.0
    @test cxx[1] == HomogeneousPolynomial(0.0,1)

    g1(x, y) = x^3 + 3y^2 - 2x^2 * y - 7x + 2
    g2(x, y) = y + x^2 - x^4
    f1 = g1(xT, yT)
    f2 = g2(xT, yT)
    @test TaylorSeries.gradient(f1) == [ 3*xT^2-4*xT*yT-TaylorN(7,0), 6*yT-2*xT^2 ]
    @test ∇(f2) == [2*xT - 4*xT^3, TaylorN(1,0)]
    @test TaylorSeries.jacobian([f1,f2], [2,1]) == TaylorSeries.jacobian( [g1(xT+2,yT+1), g2(xT+2,yT+1)] )
    jac = Array{Int}(undef, 2, 2)
    TaylorSeries.jacobian!(jac, [g1(xT+2,yT+1), g2(xT+2,yT+1)])
    @test jac == TaylorSeries.jacobian( [g1(xT+2,yT+1), g2(xT+2,yT+1)] )
    TaylorSeries.jacobian!(jac, [f1,f2], [2,1])
    @test jac == TaylorSeries.jacobian([f1,f2], [2,1])
    @test TaylorSeries.hessian( f1*f2 ) ==
        [derivative((2,0), f1*f2) derivative((1,1), (f1*f2));
         derivative((1,1), f1*f2) derivative((0,2), (f1*f2))] == [4 -7; -7 0]
    @test TaylorSeries.hessian( f1*f2, [xT, yT] ) ==
        [derivative(f1*f2, (2,0)) derivative((f1*f2), (1,1));
         derivative(f1*f2, (1,1)) derivative((f1*f2), (0,2))]
    @test [xT yT]*TaylorSeries.hessian(f1*f2)*[xT, yT] == [ 2*TaylorN((f1*f2)[2]) ]
    @test TaylorSeries.hessian(f1^2)/2 == [ [49,0] [0,12] ]
    @test TaylorSeries.hessian(f1-f2-2*f1*f2) == (TaylorSeries.hessian(f1-f2-2*f1*f2))'
    @test TaylorSeries.hessian(f1-f2,[1,-1]) == TaylorSeries.hessian(g1(xT+1,yT-1)-g2(xT+1,yT-1))
    hes = Array{Int}(undef, 2, 2)
    TaylorSeries.hessian!(hes, f1*f2)
    @test hes == TaylorSeries.hessian(f1*f2)
    @test [xT yT]*hes*[xT, yT] == [ 2*TaylorN((f1*f2)[2]) ]
    TaylorSeries.hessian!(hes, f1^2)
    @test hes/2 == [ [49,0] [0,12] ]
    TaylorSeries.hessian!(hes, f1-f2-2*f1*f2)
    @test hes == hes'
    hes1 = Array{Int}(undef, 2, 2)
    TaylorSeries.hessian!(hes1, f1-f2,[1,-1])
    TaylorSeries.hessian!(hes, g1(xT+1,yT-1)-g2(xT+1,yT-1))
    @test hes1 == hes

    use_show_default(true)
    aa = sqrt(2) * xH
    ab = sqrt(2) * TaylorN(2, order=1)
    @test string(aa) ==
        "HomogeneousPolynomial{Float64}([1.4142135623730951, 0.0], 1)"
    @test string(ab) ==
        "TaylorN{Float64}(HomogeneousPolynomial{Float64}" *
        "[HomogeneousPolynomial{Float64}([0.0], 0), " *
        "HomogeneousPolynomial{Float64}([0.0, 1.4142135623730951], 1)], 1)"
    @test string([aa, aa]) ==
        "HomogeneousPolynomial{Float64}[HomogeneousPolynomial{Float64}" *
        "([1.4142135623730951, 0.0], 1), HomogeneousPolynomial{Float64}" *
        "([1.4142135623730951, 0.0], 1)]"
    @test string([ab, ab]) == "TaylorN{Float64}[TaylorN{Float64}" *
        "(HomogeneousPolynomial{Float64}[HomogeneousPolynomial{Float64}([0.0], 0), " *
        "HomogeneousPolynomial{Float64}([0.0, 1.4142135623730951], 1)], 1), " *
        "TaylorN{Float64}(HomogeneousPolynomial{Float64}[HomogeneousPolynomial{Float64}" *
        "([0.0], 0), HomogeneousPolynomial{Float64}([0.0, 1.4142135623730951], 1)], 1)]"
    use_show_default(false)
    @test string(aa) == " 1.4142135623730951 x₁"
    @test string(ab) == " 1.4142135623730951 x₂ + 𝒪(‖x‖²)"
    displayBigO(false)
    @test string(-xH) == " - 1 x₁"
    @test string(xT^2) == " 1 x₁²"
    @test string(1im*yT) == " ( 1 im ) x₂"
    @test string(xT-im*yT) == "  ( 1 ) x₁ - ( 1 im ) x₂"
    @test string([ab, ab]) ==
        "TaylorN{Float64}[ 1.4142135623730951 x₂,  1.4142135623730951 x₂]"
    displayBigO(true)
    @test string(-xH) == " - 1 x₁"
    @test string(xT^2) == " 1 x₁² + 𝒪(‖x‖¹⁸)"
    @test string(1im*yT) == " ( 1 im ) x₂ + 𝒪(‖x‖¹⁸)"
    @test string(xT-im*yT) == "  ( 1 ) x₁ - ( 1 im ) x₂ + 𝒪(‖x‖¹⁸)"

    @test_throws ArgumentError abs(xT)
    @test_throws AssertionError 1/x
    @test_throws AssertionError zero(x)/zero(x)
    @test_throws ArgumentError sqrt(x)
    @test_throws AssertionError x^(-2)
    @test_throws ArgumentError log(x)
    @test_throws AssertionError cos(x)/sin(y)
    @test_throws BoundsError xH[20]
    @test_throws BoundsError xT[20]

    a = 3x + 4y +6x^2 + 8x*y
    @test typeof( norm(x) ) == Float64
    @test norm(x) > 0
    @test norm(a) == norm([3,4,6,8.0])
    @test norm(a, 4) == sum([3,4,6,8.0].^4)^(1/4.)
    @test norm(a, Inf) == 8.0
    @test norm((3.0 + 4im)*x) == abs(3.0 + 4im)

    @test TaylorSeries.rtoldefault(TaylorN{Int}) == 0
    @test TaylorSeries.rtoldefault(TaylorN{Float64}) == sqrt(eps(Float64))
    @test TaylorSeries.rtoldefault(TaylorN{BigFloat}) == sqrt(eps(BigFloat))
    @test TaylorSeries.real(TaylorN{Float64}) == TaylorN{Float64}
    @test TaylorSeries.real(TaylorN{Complex{Float64}}) == TaylorN{Float64}
    @test isfinite(a)
    @test a[0] ≈ a[0]
    @test a[1] ≈ a[1]
    @test a[2] ≈ a[2]
    @test a[3] ≈ a[3]
    @test a ≈ a
    @test a .≈ a
    b = deepcopy(a)
    b[2][3] = Inf
    @test !isfinite(b)
    b[2][3] = NaN
    @test !isfinite(b)
    b[2][3] = a[2][3]+eps()
    @test isapprox(a[2], b[2], rtol=eps())
    @test a ≈ b
    b[2][2] = a[2][2]+sqrt(eps())
    @test a[2] ≈ b[2]
    @test a ≈ b

    f11(a,b) = (a+b)^a - cos(a*b)*b
    f22(a) = (a[1] + a[2])^a[1] - cos(a[1]*a[2])*a[2]
    @test taylor_expand(f11, 1.0,2.0) == taylor_expand(f22, [1,2.0])
    @test evaluate(taylor_expand(x->x[1] + x[2], [1,2])) == 3.0
    f33(x,y) = 3x+y
    @test eltype(taylor_expand(f33,1,1)) == eltype(1)
    x,y = get_variables()
    xysq = x^2 + y^2
    update!(xysq,[1.0,-2.0])
    @test xysq == (x+1.0)^2 + (y-2.0)^2
    update!(xysq,[-1.0,2.0])
    @test xysq == x^2 + y^2

    #test function-like behavior for TaylorN
    @test exy() == 1
    @test exy([0.1im,0.01im]) == exp(0.11im)
    @test isapprox(exy([1,1]), eeuler^2)
    @test sin(asin(xT+yT))([1.0,0.5]) == 1.5
    @test asin(sin(xT+yT))([1.0,0.5]) == 1.5
    @test ( -sinh(xT+yT)^2 + cosh(xT+yT)^2 )(rand(2)) == 1
    @test ( -sinh(xT+yT)^2 + cosh(xT+yT)^2 )(zeros(2)) == 1
    #number of variables changed to 4...
    dx = set_variables("x", numvars=4, order=10)
    P = sin.(dx)
    v = [1.0,2,3,4]
    for i in 1:4
        @test P[i](v) == evaluate(P[i], v)
    end
    @test P.(fill(v, 4)) == fill(P(v), 4)
    F(x) = [sin(sin(x[4]+x[3])), sin(cos(x[3]-x[2])), cos(sin(x[1]^2+x[2]^2)), cos(cos(x[2]*x[3]))]
    Q = F(v+dx)
    @test Q.( fill(v, 4) ) == fill(Q(v), 4)
    vr = map(x->rand(4), 1:4)
    @test Q.(vr) == map(x->Q(x), vr)
    for i in 1:4
        @test P[i]() == evaluate(P[i])
        @test Q[i]() == evaluate(Q[i])
    end
    @test P() == evaluate.(P)
    @test P() == evaluate(P)
    @test Q() == evaluate.(Q)
    @test Q() == evaluate(Q)
    @test Q[1:3]() == evaluate(Q[1:3])

    dx = set_variables("x", numvars=4, order=10)
    for i in 1:4
        @test deg2rad(180+dx[i]) == pi + deg2rad(1.0)dx[i]
        @test rad2deg(pi+dx[i]) == 180.0+rad2deg(1.0)dx[i]
    end
    p = sin(exp(dx[1]*dx[2])+dx[3]*dx[2])/(1.0+dx[4]^2)
    q = zero(p)
    TaylorSeries.deg2rad!(q, p, 0)
    @test q[0] == p[0]*(pi/180)
    # TaylorSeries.deg2rad!.(q, p, [1,3,5])
    # for i in [0,1,3,5]
    #     @test q[i] == p[i]*(pi/180)
    # end
    TaylorSeries.rad2deg!(q, p, 0)
    @test q[0] == p[0]*(180/pi)
    # TaylorSeries.rad2deg!.(q, p, [1,3,5])
    # for i in [0,1,3,5]
    #     @test q[i] == p[i]*(180/pi)
    # end
    xT = 5+TaylorN(Int, 1, order=10)
    yT = TaylorN(2, order=10)
    TaylorSeries.deg2rad!(yT, xT, 0)
    @test yT[0] == xT[0]*(pi/180)
    TaylorSeries.rad2deg!(yT, xT, 0)
    @test yT[0] == xT[0]*(180/pi)
end

@testset "Integrate for several variables" begin

    t, x, y = set_variables("t x y")

    @test integrate(t, 1) == 0.5*t^2
    @test integrate(t, 2) == t * x
    @test integrate(t, 3) == t * y
    @test integrate(x, 1) == t * x
    @test integrate(x, 2) == 0.5*x^2
    @test integrate(y, 2) == x * y

end


"""
Print the sequence of outbound nodes of a graph starting from the given index.
pystyle implies reversing the final order and substracting 1 from the vertex number.
"""
function prettyp(G::AbstractGraph, strtidx::Int; pystyle=false)
    arrow = G isa LightGraphs.SimpleDiGraph ? "->" : "--"
    v = unroll(G, strtidx)
    println(join(string.(pystyle ? reverse(v .- 1) : v), arrow))
end

"""
Walk the given graph, starting from the given vertex,
returning the list of all outbound vertices that are visited.
"""
function unroll(G::AbstractGraph, idx::Int; visited=Int[])
    idx in visited && return visited
    push!(visited, idx)
    N = outneighbors(G, idx)
    length(N) < 1 && length(visited) == nv(G) ? (return visited) : nothing
    for n in N
        unroll(G, n, visited=visited)
    end
    return reverse(visited)
end

"""
Return the first and last `len` elements of the vector as a string.

Default is to print 3 elements at head and tail.

```julia
julia> v = collect(1:10);
julia> prettyp(v)

"1,2,3...8,9,10"
```

Example w/5 elements visible

```julia
julia> v = collect(1:10);
julia> prettyp(v, 5)

"1,2,3,4,5...6,7,8,9,10"
```

"""
function prettyp(v::AbstractVector, len::Int = 3)
    length(v) < 7 && return join(string.(v), ",")
    head = join(string.(v[1:len]),",")
    tail = join(string.(v[(end-(len-1)):end]),",")
    return join([head,tail],"...")
end


# ------------------------------------------------------------------------------------------
# # 多重派发
#
# **多重派发**是这篇notebook要探索的Julia的一个关键特性。
#
# 它有助于快速开发。它也让软件可拓展、可设计和更好把玩。
#
# 它可能是并行计算取得重大突破的预兆。
#
# ## 大纲
# 1. 罗马数字
# 2. 函数
# 3. 并行计算
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ## 1. 罗马数字（玩玩）
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# 我们来定义一个**新结构**表示罗马数字。为了方便写代码，我们将只考虑 0 到 9。
#
# **练习**：拓展到更大的数字。（别忘了罗马数字是十进制计数系统！）
# ------------------------------------------------------------------------------------------

struct Roman
    n::Int
end

Base.show(io::IO, r::Roman) = print(io, 'ⅰ' + (r.n - 1) % 10 )  # 漂亮的显示；'ⅰ'是一个Unicode罗马数字

# ------------------------------------------------------------------------------------------
# 我们可以像这样创建一个这个类型的对象：
# ------------------------------------------------------------------------------------------

Roman(4)

typeof.([5 5.0 Roman(5) "Five" '5'  5//1])

# ------------------------------------------------------------------------------------------
# 用罗马数字漂亮地打印数组：
# ------------------------------------------------------------------------------------------

x = [7 1 2 5 8 9]
Roman.(x)   # equivalent to map(Roman, x)  or  [Roman(w) for w in x]

# ------------------------------------------------------------------------------------------
# 要是能像普通数字那样给罗马数字做加法就好了：
# ------------------------------------------------------------------------------------------

Roman(4) + Roman(5)

# ------------------------------------------------------------------------------------------
# 但 Julia 不知道该怎么做。
# 我们可以通过 `import` 函数 `+` 并拓展它的定义来教 Julia 给罗马数字做加法：
# ------------------------------------------------------------------------------------------

import Base: +, *

+(a::Roman, b::Roman) = Roman(a.n + b.n)

Roman(4) + Roman(5)

# ------------------------------------------------------------------------------------------
# 这样就给函数 `+` **添加了一个新方法**：
# ------------------------------------------------------------------------------------------

methods(+)

@which Roman(4) + Roman(5)

import Base.*
*(i::Roman, j::Roman) = Roman(i.n * j.n)                     # Multiply like a Roman

Roman(3) * Roman(2)

Roman.(1:3) .* [Roman(1) Roman(2) Roman(3)]

# ------------------------------------------------------------------------------------------
# 但是我们的乘法仍有一些问题
# ------------------------------------------------------------------------------------------

Roman(3) * 2

# mytimes 函数目前基于类型来判断该做什么，很复杂
# 不用担心，更好的方法马上就来
function mytimes(i,j)
  if isa(i,Roman) & isa(j,Number)
        return  fill(1, i.n, j)   # i by j matrix with ones
    elseif    isa(i,Number) & isa(j,Roman)
        return "😄"^ (i*j.n)   #  i * j happy faces
    else
        return("I Don't know")
    end
end

mytimes(4,Roman(3)) # 12个笑脸

mytimes(Roman(4),3) # 4x3的全1矩阵

# ------------------------------------------------------------------------------------------
# 最简单的实现方式是明确定义一个 `Roman` 和一个数字的乘法。我们可以按我们的想法来定义：
# ------------------------------------------------------------------------------------------

*(i::Number, j::Roman) = "😄"^ (i*j.n)    # i * j个笑脸

*(i::Roman, j::Number) =   fill(1, i.n, j) # i * j的矩阵

3 * Roman(3) # 9个笑脸

Roman(3) * 5  # 3*5的全1矩阵

t(x::Roman,y::Roman) = x.n * y.n

t(Roman(5),Roman(4))

# 注意它的汇编代码是多紧凑！
@code_native t(Roman(2),Roman(4))

# ------------------------------------------------------------------------------------------
# ## 函数
# ------------------------------------------------------------------------------------------

import Base: *, +, ^

*(α::Number,   g::Function) = x -> α * g(x)   # 标量乘以函数

*(f::Function, λ::Number)   = x -> f(λ * x)   # Scale the argument

*(f::Function, g::Function) = x -> f(g(x))    # 复合函数  -- 滥用符号！  在Julia0.6中使用 \circ

^(f::Function, n::Integer) = n == 1 ? f : f*f^(n-1) # 一个通过递归乘法实现的天真的求幂算法

+(f::Function, g::Function) = x -> f(x) + g(x)

# ------------------------------------------------------------------------------------------
# 举个例子，定义成这样的指数函数
#
# $$\exp(x) = \sum_{n=0}^\infty \frac{1}{n!} x^n.$$
#
# 我们可以把它看成是这样的函数：
#
# $$\exp = \sum_{n=0}^\infty \frac{1}{n!} \mathrm{pow}_n,$$
#
# 其中 $\mathrm{pow}_n(x) = x^n$.
#
# （开始用数字模糊符号！）
# ------------------------------------------------------------------------------------------

pow(n) = x -> x^n

myexp = sum(1/factorial(big(n)) * pow(n) for n in 0:100)   # 效率低的泰勒级数！

[myexp(1); exp(1); exp(big(1))]

f = x -> x^2
f(10)

g = 3f
g(10)

(f^2)(10)  # 因为我们已经定义了函数乘法为符合函数

using Plots;
gr()

x = pi*(0:0.001:4)

plot( x,  sin.(x),        c="black",            label="Fun")
plot!(x, (12*sin).(x),    c="green",            label="Num * Fun")
plot!(x, (sin*12).(x),    c="red",   alpha=0.9, label="Fun * Num")
plot!(x, (5*sin*exp).(x), c="blue",  alpha=0.2, label="Num * Fun * Fun")

plot([12*sin, sin*12, 5*sin*exp], 0:.01:4π, α=[1 .9 .2], c=[:green :red :blue])

# ------------------------------------------------------------------------------------------
# > “我很讨厌 $sin^2 \phi$，虽然 Laplace 用过它；
# > 应该要担心 $sin^2 \phi$ 可能会引起歧义，如果说 $sin(\phi^2)$ 就不会或者说几乎不会引起歧义了，那么我们应该写 $(sin \phi)^2$，而不是
# $sin^2 \phi$，以此类推 $sin^2 \phi$应该指的是 $sin(sin \phi)$。”
# >
# > —— Gauss
# ------------------------------------------------------------------------------------------

x=(0:.01:2) * pi;

plot(x, (sin^2).(x), c="blue")     # 乘方能工作，y=sin(sin(x))，Gauss会开心的！
plot!(x, sin.(x).^2, c="red")

# ------------------------------------------------------------------------------------------
# ## 练习
# ------------------------------------------------------------------------------------------

h(a, b::Any) = "fallback"
h(a::Number, b::Number) = "a and b are both numbers"
h(a::Number, b) = "a is a number"
h(a, b::Number) = "b is a number"
h(a::Integer, b::Integer) = "a and b are both integers"

# 试着把玩 h，尝试使用 h 的 5 种方法
# 在此作答


function Smat(n, lam)
    if lam == 0
        e = ones(n)
        S = myspdiagm(spdiagm(0 => 0.5*e, 2 => -0.5*e), n)
        S[1] = 1
    else
        e = lam./(collect(0:n-1).+lam)
        S = myspdiagm(spdiagm(0 => e, 2 => -e), n)
    end
    return S
end

myspdiagm(M, N) = size(M, 1) !== N ? (return M[1:N, 1:N]) : (return M)


import Base: -,*, convert
using LinearAlgebra

export UnitaryGate, tracenorm, tracedist, average_fidelity, dim, ⊗

abstract type QuantumGate end

struct UnitaryGate{T <: Complex} <: QuantumGate
    U::Matrix{T}
end

function UnitaryGate(U::T) where T<:AbstractMatrix
    if !isunitary(U)
        error("Gate must be unitary!")
    end
    (n,m) = size(U)
    if !(n == m)
        error("Gate matrix must be square!")
    end
    UnitaryGate(complex.(U))
end

Matrix(G::UnitaryGate{T}) where {T} = G.U

*(A::T, B::T) where T <: UnitaryGate = UnitaryGate(A.U*B.U)
*(A::T, B::AbstractMatrix) where T<:UnitaryGate = UnitaryGate(A.U*B)
*(A::AbstractMatrix, B::T) where T<:UnitaryGate = UnitaryGate(A*B.U)
*(A::Tg, x::Tx) where {Tg<:UnitaryGate, Tx<:Number} = UnitaryGate(A.U*x)
*(x::Tx, A::Tg) where {Tg<:UnitaryGate, Tx<:Number} = UnitaryGate(x*A.U)
-(A::T, B::T) where T<: UnitaryGate = UnitaryGate(A.U-B.U)

LinearAlgebra.kron(A::T, B::T) where T<:UnitaryGate = UnitaryGate(kron(A.U, B.U))
⊗(A::T, B::T) where T<:UnitaryGate = kron(A, B)

LinearAlgebra.adjoint(A::UnitaryGate) = UnitaryGate(Matrix(A.U'))
LinearAlgebra.tr(A::UnitaryGate) = tr(A.U)

dim(G::T) where T<:UnitaryGate = size(G.U)[1]

function dim(U::T)::Float64 where T<:AbstractMatrix
    if size(U)[1] == size(U)[2]
        return size(U)[1]
    else
        error("Can only return dimension of square matrix!")
    end
end

tracenorm(G::UnitaryGate) = tracenorm(Matrix(G))
tracedist(A::UnitaryGate, B::UnitaryGate)  = tracedist(Matrix(A), Matrix(B))
frobdist(A::UnitaryGate, B::UnitaryGate) = frobdist(Matrix(A), Matrix(B))

LinearAlgebra.norm(G::UnitaryGate) = norm(G.U, 2)
LinearAlgebra.ishermitian(G::T; tol=ϵ) where T<:UnitaryGate = ishermitian_tol(Matrix(G), tol=ϵ)

"""
Returns the average gate fidelity between two gates.
"""
function average_fidelity(A::T, B::T) where T<:AbstractMatrix
    @assert dim(A) == dim(B)
    d = dim(A)
    return (d + abs(tr(A*B'))^2)/(d*(d+1))
end

average_fidelity(A::T, B::T) where T <: UnitaryGate = average_fidelity(A.U, B.U)


### A Pluto.jl notebook ###
# v0.12.21

using Markdown
using InteractiveUtils

# ╔═╡ 102adb52-8b83-11eb-1876-a52236a70e1f
using Images,MosaicViews

# ╔═╡ 4f027e0c-8b88-11eb-2f3f-237e69a18582
begin
	using TestImages, ImageContrastAdjustment, HistogramThresholding

img = testimage("cameraman")
end

# ╔═╡ 5bd1efe0-8bf2-11eb-3f4c-6d5a43febc6d
begin
edges, counts = build_histogram(img,256)
t = find_threshold(Otsu(), counts[1:end], edges)
# Apply the threshold to img and visualize the result.
img_final = zeros(Gray{Bool}, axes(img))
for i in CartesianIndices(img)
        img_final[i] = img[i] < t ? 0 : 1
end
img_final
end

# ╔═╡ Cell order:
# ╠═102adb52-8b83-11eb-1876-a52236a70e1f
# ╠═4f027e0c-8b88-11eb-2f3f-237e69a18582
# ╠═5bd1efe0-8bf2-11eb-3f4c-6d5a43febc6d


using Test
using BSeries

using BSeries.Latexify: latexify

using StaticArrays: @SArray, @SMatrix, @SVector

using SymEngine: SymEngine
using SymPy: SymPy
using Symbolics: Symbolics

@testset "BSeries" begin

@testset "lazy representation of exact ODE solution" begin
  exact = ExactSolution{Rational{Int}}()
  terms = collect(Iterators.take(exact, 4))
  @test terms == [1//1, 1//2, 1//6, 1//3]
  @test exact == ExactSolution(exact)
end

@testset "latexify" begin
  # explicit midpoint method
  A = @SArray [0 0; 1//2 0]; b = @SArray [0, 1//1]; c = @SArray [0, 1//2];

  series_integrator = bseries(A, b, c, 2)
  @test_nowarn latexify(series_integrator)
  @test_nowarn latexify(series_integrator, cdot=false)
  @test_nowarn latexify(series_integrator, dt=SymEngine.symbols("h"))

  @testset "SymEngine.jl" begin
    α = SymEngine.symbols("α")
    A = [0 0; 1/(2*α) 0]; b = [1-α, α]; c = [0, 1/(2*α)]
    series_integrator = bseries(A, b, c, 3)
    @test_nowarn latexify(series_integrator)
  end

  @testset "SymPy.jl" begin
    α = SymPy.symbols("α", real=true)
    A = [0 0; 1/(2*α) 0]; b = [1-α, α]; c = [0, 1/(2*α)]
    series_integrator = bseries(A, b, c, 3)
    @test_nowarn latexify(series_integrator)
  end

  @testset "Symbolics.jl" begin
    Symbolics.@variables α
    A = [0 0; 1/(2*α) 0]; b = [1-α, α]; c = [0, 1/(2*α)]
    series_integrator = bseries(A, b, c, 3)
    @test_nowarn latexify(series_integrator)
  end
end

@testset "AbstractDict interface" begin
  # explicit midpoint method
  A = @SArray [0 0; 1//2 0]; b = @SArray [0, 1//1]; c = @SArray [0, 1//2];
  series_integrator = bseries(A, b, c, 2)

  # These are mostly simple smoke tests
  @test_nowarn begin
    @test isempty(empty(series_integrator, RootedTrees.RootedTree{Int32, Vector{Int32}}, Rational{Int32}))
  end
  empty!(series_integrator)
  @test isempty(series_integrator)
  @test isempty(keys(series_integrator))
  @test isempty(values(series_integrator))
  @test length(series_integrator) == 0
  @test get(series_integrator, rootedtree([1]), :default) == :default
  @test get(series_integrator, rootedtree([1])) do
    :default
  end == :default
  @test getkey(series_integrator, rootedtree([1]), :default) == :default
  @test_nowarn sizehint!(series_integrator, 10)

  @test get!(series_integrator, rootedtree([1]), 2) == 2
  @test pop!(series_integrator, rootedtree([1])) == 2
  @test pop!(series_integrator, rootedtree([1]), :default) == :default
  series_integrator[rootedtree([1])] = 3
  @test isempty(delete!(series_integrator, rootedtree([1])))
end

@testset "substitute" begin
  Symbolics.@variables a1 a2 a31 a32
  a = OrderedDict{RootedTrees.RootedTree{Int, Vector{Int}}, Symbolics.Num}(
        rootedtree(Int[])     => 1,
        rootedtree([1])       => a1,
        rootedtree([1, 2])    => a2,
        rootedtree([1, 2, 2]) => a31,
        rootedtree([1, 2, 3]) => a32)

  Symbolics.@variables b1 b2 b31 b32
  b = OrderedDict{RootedTrees.RootedTree{Int, Vector{Int}}, Symbolics.Num}(
        rootedtree(Int[])     => 0,
        rootedtree([1])       => b1,
        rootedtree([1, 2])    => b2,
        rootedtree([1, 2, 2]) => b31,
        rootedtree([1, 2, 3]) => b32)

  # See equation (6) of
  # - Philippe Chartier, Ernst Hairer and Gilles Vilmart (2007)
  #   Numerical integrators based on modified differential equations
  #   [DOI: 10.1090/S0025-5718-07-01967-9](https://doi.org/10.1090/S0025-5718-07-01967-9)

  ba = @inferred substitute(b, a)
  t = rootedtree(Int[])
  @test isequal(ba[t], a[t])

  t = rootedtree([1])
  coef = a1 * b1
  @inferred substitute(b, a, t)
  @test isequal(substitute(b, a, t), coef)
  @test isequal(ba[t], coef)

  t = rootedtree([1, 2])
  coef = a1 * b2 + a2 * b1^2
  @inferred substitute(b, a, t)
  @test isequal(substitute(b, a, t), coef)
  @test isequal(ba[t], coef)

  t = rootedtree([1, 2, 2])
  coef = a1 * b31 + 2 * a2 * b1 * b2 + a31 * b1^3
  @inferred substitute(b, a, t)
  @test isequal(substitute(b, a, t), coef)
  @test isequal(ba[t], coef)

  t = rootedtree([1, 2, 3])
  coef = a1 * b32 + 2 * a2 * b1 * b2 + a32 * b1^3
  @inferred substitute(b, a, t)
  @test isequal(substitute(b, a, t), coef)
  @test isequal(ba[t], coef)
end

@testset "compose" begin
  Symbolics.@variables a0 a1 a2 a31 a32
  a = OrderedDict{RootedTrees.RootedTree{Int, Vector{Int}}, Symbolics.Num}(
        rootedtree(Int[])     => a0,
        rootedtree([1])       => a1,
        rootedtree([1, 2])    => a2,
        rootedtree([1, 2, 2]) => a31,
        rootedtree([1, 2, 3]) => a32)

  Symbolics.@variables b1 b2 b31 b32
  b = OrderedDict{RootedTrees.RootedTree{Int, Vector{Int}}, Symbolics.Num}(
        rootedtree(Int[])     => 0,
        rootedtree([1])       => b1,
        rootedtree([1, 2])    => b2,
        rootedtree([1, 2, 2]) => b31,
        rootedtree([1, 2, 3]) => b32)

  # See equation (5) of
  # - Philippe Chartier, Ernst Hairer and Gilles Vilmart (2007)
  #   Numerical integrators based on modified differential equations
  #   [DOI: 10.1090/S0025-5718-07-01967-9](https://doi.org/10.1090/S0025-5718-07-01967-9)
  t = rootedtree([1])
  @inferred compose(b, a, t)
  @test isequal(compose(b, a, t), a0 * b1 + a1)

  t = rootedtree([1, 2])
  @inferred compose(b, a, t)
  @test isequal(compose(b, a, t), a0 * b2 + a1 * b1 + a2)

  t = rootedtree([1, 2, 2])
  @inferred compose(b, a, t)
  @test isequal(compose(b, a, t), a0 * b31 + a1 * b1^2 + 2 * a2 * b1 + a31)

  t = rootedtree([1, 2, 3])
  @inferred compose(b, a, t)
  @test isequal(compose(b, a, t), a0 * b32 + a1 * b2 + a2 * b1 + a32)

  @testset "Composing RK4 with itself" begin
    # classical fourth-order Runge-Kutta method
    A = @SArray [0 0 0 0;
                 1//2 0 0 0;
                 0 1//2 0 0;
                 0 0 1 0]
    b = @SArray [1//6, 1//3, 1//3, 1//6]
    c = @SArray [0, 1//2, 1//2, 1]

    series_rk4 = bseries(A, b, c, 8)
    series_rk4_composed = @inferred compose(series_rk4, series_rk4,
                                            normalize_stepsize=true)

    # Butcher coefficients of the composition (with normalized step size)
    A = @SArray [0     0     0     0     0     0     0     0;
                 1//4  0     0     0     0     0     0     0;
                 0     1//4  0     0     0     0     0     0;
                 0     0     1//2  0     0     0     0     0;
                 1//12 1//6  1//6  1//12 0     0     0     0;
                 1//12 1//6  1//6  1//12 1//4  0     0     0;
                 1//12 1//6  1//6  1//12 0     1//4  0     0;
                 1//12 1//6  1//6  1//12 0     0     1//2  0]
    b = @SArray [1//12, 1//6, 1//6, 1//12, 1//12, 1//6, 1//6, 1//12]
    c = @SArray [0, 1//4, 1//4, 1//2, 1//2, 3//4, 3//4, 1//1]

    series_2rk4 = bseries(A, b, c, order(series_rk4))

    @test series_rk4_composed == series_2rk4

    # Ensure that the normalization uses the correct factor
    unnormalized_series = @inferred compose(series_rk4, series_rk4,
                                            normalize_stepsize=false)
    for t in keys(unnormalized_series)
      @test unnormalized_series[t] == series_rk4_composed[t] * 2^order(t)
    end
  end
end

@testset "modified_equation" begin
  t1  = rootedtree([1])
  t2  = rootedtree([1, 2])
  t31 = rootedtree([1, 2, 2])
  t32 = rootedtree([1, 2, 3])
  t41 = rootedtree([1, 2, 2, 2])
  t42 = rootedtree([1, 2, 2, 3])
  t43 = rootedtree([1, 2, 3, 3])
  t44 = rootedtree([1, 2, 3, 4])
  t51 = rootedtree([1, 2, 2, 2, 2])
  t52 = rootedtree([1, 2, 2, 2, 3])
  t53 = rootedtree([1, 2, 2, 3, 3])
  t54 = rootedtree([1, 2, 2, 3, 4])
  t55 = rootedtree([1, 2, 3, 2, 3])
  t56 = rootedtree([1, 2, 3, 3, 3])
  t57 = rootedtree([1, 2, 3, 3, 4])
  t58 = rootedtree([1, 2, 3, 4, 4])
  t59 = rootedtree([1, 2, 3, 4, 5])

  @testset "SSPRK(2, 2) aka Heun's method" begin
    A = @SArray [0.0 0.0;
                1.0 0.0]
    b = @SArray [1/2, 1/2]
    c = @SArray [0.0, 1.0]
    order = 5
    series = modified_equation(A, b, c, order)

    # tested with the Python package BSeries
    @test series[t1 ] ≈ 1.0                   atol=10*eps()
    @test series[t2 ] ≈ 0                     atol=10*eps()
    @test series[t31] ≈ 0.166666666666667     atol=10*eps()
    @test series[t32] ≈ -0.166666666666667    atol=10*eps()
    @test series[t41] ≈ -1.11022302462516e-16 atol=10*eps()
    @test series[t42] ≈ -0.125000000000000    atol=10*eps()
    @test series[t43] ≈ -2.77555756156289e-17 atol=10*eps()
    @test series[t44] ≈ 0.125000000000000     atol=10*eps()
    @test series[t51] ≈ -0.0333333333333332   atol=10*eps()
    @test series[t52] ≈ -0.0583333333333333   atol=10*eps()
    @test series[t55] ≈ 0.0750000000000000    atol=10*eps()
    @test series[t53] ≈ 0.0583333333333333    atol=10*eps()
    @test series[t56] ≈ 0.0333333333333334    atol=10*eps()
    @test series[t54] ≈ 0.0500000000000000    atol=10*eps()
    @test series[t57] ≈ 0.0583333333333333    atol=10*eps()
    @test series[t58] ≈ -0.0583333333333333   atol=10*eps()
    @test series[t59] ≈ -0.0500000000000000   atol=10*eps()
  end

  @testset "SSPRK(3, 3)" begin
    A = @SArray [0.0 0.0 0.0;
                1.0 0.0 0.0;
                1/4 1/4 0.0]
    b = @SArray [1/6, 1/6, 2/3]
    c = @SArray [0.0, 1.0, 1/2]
    order = 5
    series = modified_equation(A, b, c, order)

    # tested with the Python package BSeries
    @test series[t1 ] ≈ 1.0                  atol=10*eps()
    @test series[t2 ] ≈ 0                    atol=10*eps()
    @test series[t31] ≈ 0                    atol=10*eps()
    @test series[t32] ≈ 0                    atol=10*eps()
    @test series[t41] ≈ 0                    atol=10*eps()
    @test series[t42] ≈ -0.0416666666666667  atol=10*eps()
    @test series[t43] ≈ 0.0833333333333333   atol=10*eps()
    @test series[t44] ≈ -0.0416666666666667  atol=10*eps()
    @test series[t51] ≈ 0.00833333333333330  atol=10*eps()
    @test series[t52] ≈ -0.0166666666666667  atol=10*eps()
    @test series[t55] ≈ 0.0333333333333333   atol=10*eps()
    @test series[t53] ≈ 0.0166666666666667   atol=10*eps()
    @test series[t56] ≈ -0.00833333333333333 atol=10*eps()
    @test series[t54] ≈ 0.00833333333333333  atol=10*eps()
    @test series[t57] ≈ -0.0250000000000000  atol=10*eps()
    @test series[t58] ≈ -0.0166666666666667  atol=10*eps()
    @test series[t59] ≈ 0.0333333333333333   atol=10*eps()
  end

  @testset "Explicit midpoint method (Runge's method)" begin
    A = @SArray [0 0; 1//2 0]
    b = @SArray [0, 1//1]
    c = @SArray [0, 1//2]
    order = 5
    series = modified_equation(A, b, c, order)

    # tested with the Python package BSeries
    @test series[t1 ] == 1
    @test series[t2 ] == 0
    @test series[t31] == -1//12
    @test series[t32] == -1//6
    @test series[t41] == 0
    @test series[t42] == 0
    @test series[t43] == 1//8
    @test series[t44] == 1//8
    @test series[t51] == 7//240
    @test series[t52] == 1//40
    @test series[t55] == 1//30
    @test series[t53] == 3//80
    @test series[t56] == -7//240
    @test series[t54] == 7//240
    @test series[t57] == -1//40
    @test series[t58] == -19//240
    @test series[t59] == -1//20
  end
end

@testset "elementary differentials" begin
  @testset "Bicolored trees" begin
    @testset "Lotka-Volterra" begin
      # Verified with Mathematica
      # u = {q, p};
      # f1 = {q * (p - 1), 0};
      # f2 = {0, p * (2 - q)};
      #
      # f1p = Simplify@D[f1, {u}];
      # f2p = Simplify@D[f2, {u}];
      #
      # f1pp = Simplify@D[f1, {u, 2}];
      # f2pp = Simplify@D[f2, {u, 2}];
      #
      # (* Trees of order 2 *)
      # Map[InputForm, f1p.f1]
      # Map[InputForm, f2p.f1]
      # Map[InputForm, f1p.f2]
      # Map[InputForm, f2p.f2]
      #
      # (* Tall trees of order 3 *)
      # Map[InputForm, f1p.(f1p.f1)]
      # Map[InputForm, f2p.(f1p.f1)]
      # Map[InputForm, f1p.(f2p.f1)]
      # Map[InputForm, f2p.(f2p.f1)]
      # Map[InputForm, f1p.(f1p.f2)]
      # Map[InputForm, f2p.(f1p.f2)]
      # Map[InputForm, f1p.(f2p.f2)]
      # Map[InputForm, f2p.(f2p.f2)]
      #
      # (* Bushy trees of order 3 *)
      # Map[InputForm, f1pp.f1.f1]
      # Map[InputForm, f2pp.f1.f1]
      # Map[InputForm, f1pp.f2.f1]
      # Map[InputForm, f2pp.f2.f1]
      # (* The other choice of colors of the leaves is redundant *)
      # Map[InputForm, f1pp.f2.f2]
      # Map[InputForm, f2pp.f2.f2]
      @testset "SymEngine" begin
        q, p = u = SymEngine.symbols("q, p")
        f = ([q * (p - 1), 0], [0, p * (2 - q)])
        differentials = elementary_differentials(f, u, 3)
        @test differentials[rootedtree(Int64[], Bool[])         ] == [1, 1]
        @test differentials[rootedtree([1], Bool[0])            ] == [q * (p - 1), 0]
        @test differentials[rootedtree([1], Bool[1])            ] == [0, p * (2 - q)]
        @test differentials[rootedtree([1, 2], Bool[0, 0])      ] == [(-1 + p)^2*q,0]
        @test differentials[rootedtree([1, 2], Bool[1, 0])      ] == [0,-((-1 + p)*p*q)]
        @test differentials[rootedtree([1, 2], Bool[0, 1])      ] == [p*(2 - q)*q,0]
        @test differentials[rootedtree([1, 2], Bool[1, 1])      ] == [0,p*(2 - q)^2]
        @test differentials[rootedtree([1, 2, 3], Bool[0, 0, 0])] == [(-1 + p)^3*q,0]
        @test differentials[rootedtree([1, 2, 3], Bool[1, 0, 0])] == [0,-((-1 + p)^2*p*q)]
        @test differentials[rootedtree([1, 2, 3], Bool[0, 1, 0])] == [-((-1 + p)*p*q^2),0]
        @test differentials[rootedtree([1, 2, 3], Bool[1, 1, 0])] == [0,-((-1 + p)*p*(2 - q)*q)]
        @test differentials[rootedtree([1, 2, 3], Bool[0, 0, 1])] == [(-1 + p)*p*(2 - q)*q,0]
        @test differentials[rootedtree([1, 2, 3], Bool[1, 0, 1])] == [0,-(p^2*(2 - q)*q)]
        @test differentials[rootedtree([1, 2, 3], Bool[0, 1, 1])] == [p*(2 - q)^2*q,0]
        @test differentials[rootedtree([1, 2, 3], Bool[1, 1, 1])] == [0,p*(2 - q)^3]
        @test differentials[rootedtree([1, 2, 2], Bool[0, 0, 0])] == [0, 0]
        @test differentials[rootedtree([1, 2, 2], Bool[1, 0, 0])] == [0, 0]
        @test differentials[rootedtree([1, 2, 2], Bool[0, 1, 0])] == [(-1 + p)*p*(2 - q)*q,0]
        @test differentials[rootedtree([1, 2, 2], Bool[1, 1, 0])] == [0,-((-1 + p)*p*(2 - q)*q)]
        @test differentials[rootedtree([1, 2, 2], Bool[0, 1, 1])] == [0, 0]
        @test differentials[rootedtree([1, 2, 2], Bool[1, 1, 1])] == [0, 0]
      end
    end
  end
end

@testset "modified_equation with elementary differentials" begin
  @testset "SymEngine.jl" begin
    @testset "Explicit Euler" begin
      # Lotka-Volterra model
      dt = SymEngine.symbols("dt")
      p, q = u = SymEngine.symbols("p, q")
      f = [p * (2 - q), q * (p - 1)]

      # Explicit Euler method
      A = @SArray [0//1;]
      b = @SArray [1//1]
      c = @SArray [0//1]

      # tested with the Python package BSeries
      series2 = modified_equation(f, u, dt, A, b, c, 2)
      series2_reference = [
        -dt*(-p*q*(p - 1) + p*(2 - q)^2)/2 + p*(2 - q),
        -dt*( p*q*(2 - q) + q*(p - 1)^2)/2 + q*(p - 1)
      ]
      @test mapreduce(isequal, &, series2, series2_reference)

      # tested with the Python package BSeries
      series3 = modified_equation(f, u, dt, A, b, c, 3)
      series3_reference = [
        -dt^2*p*q*(2 - q)*(p - 1)/6 + dt^2*(-p^2*q*(2 - q) - p*q*(2 - q)*(p - 1) - p*q*(p - 1)^2 + p*(2 - q)^3)/3 - dt*(-p*q*(p - 1) + p*(2 - q)^2)/2 + p*(2 - q),
        dt^2*p*q*(2 - q)*(p - 1)/6 + dt^2*(-p*q^2*(p - 1) + p*q*(2 - q)^2 + p*q*(2 - q)*(p - 1) + q*(p - 1)^3)/3 - dt*(p*q*(2 - q) + q*(p - 1)^2)/2 + q*(p - 1)
      ]
      @test mapreduce(iszero ∘ SymEngine.expand, &, series3 - series3_reference)

      # tested with the Python package BSeries
      series4 = modified_equation(f, u, dt, A, b, c, 4)
      series4_reference = [
        -dt^3*(-2*p^2*q*(2 - q)*(p - 1) + 2*p*q*(2 - q)*(p - 1)*(q - 2))/12 - dt^3*(-p^2*q*(2 - q)^2 + p*q^2*(p - 1)^2 + p*q*(1 - p)*(2 - q)*(p - 1) + p*q*(2 - q)*(p - 1)*(q - 2))/12 - dt^3*(p^2*q^2*(p - 1) - 2*p^2*q*(2 - q)^2 - p^2*q*(2 - q)*(p - 1) - p*q*(2 - q)^2*(p - 1) - p*q*(2 - q)*(p - 1)^2 - p*q*(p - 1)^3 + p*(2 - q)^4)/4 - dt^2*p*q*(2 - q)*(p - 1)/6 + dt^2*(-p^2*q*(2 - q) - p*q*(2 - q)*(p - 1) - p*q*(p - 1)^2 + p*(2 - q)^3)/3 - dt*(-p*q*(p - 1) + p*(2 - q)^2)/2 + p*(2 - q),
          -dt^3*(-2*p*q^2*(2 - q)*(p - 1) + 2*p*q*(2 - q)*(p - 1)^2)/12 - dt^3*(p^2*q*(2 - q)^2 - p*q^2*(p - 1)^2 + p*q*(2 - q)^2*(p - 1) + p*q*(2 - q)*(p - 1)^2)/12 - dt^3*(-p^2*q^2*(2 - q) - p*q^2*(2 - q)*(p - 1) - 2*p*q^2*(p - 1)^2 + p*q*(2 - q)^3 + p*q*(2 - q)^2*(p - 1) + p*q*(2 - q)*(p - 1)^2 + q*(p - 1)^4)/4 + dt^2*p*q*(2 - q)*(p - 1)/6 + dt^2*(-p*q^2*(p - 1) + p*q*(2 - q)^2 + p*q*(2 - q)*(p - 1) + q*(p - 1)^3)/3 - dt*(p*q*(2 - q) + q*(p - 1)^2)/2 + q*(p - 1)
      ]
      @test mapreduce(iszero ∘ SymEngine.expand, &, series4 - series4_reference)
    end

    @testset "IMEX Euler (partitioned Euler, symplectic Euler)" begin
      # Lotka-Volterra model
      dt = SymEngine.symbols("dt")
      q, p = u = SymEngine.symbols("q, p")
      f = ([0, p * (2 - q)], [q * (p - 1), 0])

      # Symplectic Euler method
      ex_euler = RungeKuttaMethod(
        @SMatrix([0]), @SVector [1]
      )
      im_euler = RungeKuttaMethod(
        @SMatrix([1]), @SVector [1]
      )
      ark = AdditiveRungeKuttaMethod([im_euler, ex_euler])

      # Hairer, Lubich, Wanner (2006) Geometric numerical integration
      # Example IX.1.3(b)
      series_integrator = bseries(ark, 2)
      series = modified_equation(f, u, dt, series_integrator)
      series_reference = [
        q*(p - 1) - dt / 2 * q * (p^2 + p * q - 4 * p + 1),
        -p * (q - 2) + dt / 2 * p * (q^2 + p * q - 5 * q + 4)
      ]
      @test mapreduce(iszero ∘ SymEngine.expand, &, series - series_reference)
    end
  end

  @testset "SymPy.jl" begin
    # Lotka-Volterra model
    dt = SymPy.symbols("dt")
    p, q = u = SymPy.symbols("p, q")
    f = [p * (2 - q), q * (p - 1)]

    # Explicit Euler method
    A = @SArray [0//1;]
    b = @SArray [1//1]
    c = @SArray [0//1]

    # tested with the Python package BSeries
    series2 = modified_equation(f, u, dt, A, b, c, 2)
    series2_reference = [
      -dt*(-p*q*(p - 1) + p*(2 - q)^2)/2 + p*(2 - q),
      -dt*( p*q*(2 - q) + q*(p - 1)^2)/2 + q*(p - 1)
    ]
    @test mapreduce(isequal, &, series2, series2_reference)

    # tested with the Python package BSeries
    series3 = modified_equation(f, u, dt, A, b, c, 3)
    series3_reference = [
      -dt^2*p*q*(2 - q)*(p - 1)/6 + dt^2*(-p^2*q*(2 - q) - p*q*(2 - q)*(p - 1) - p*q*(p - 1)^2 + p*(2 - q)^3)/3 - dt*(-p*q*(p - 1) + p*(2 - q)^2)/2 + p*(2 - q),
      dt^2*p*q*(2 - q)*(p - 1)/6 + dt^2*(-p*q^2*(p - 1) + p*q*(2 - q)^2 + p*q*(2 - q)*(p - 1) + q*(p - 1)^3)/3 - dt*(p*q*(2 - q) + q*(p - 1)^2)/2 + q*(p - 1)
    ]
    @test mapreduce(iszero ∘ SymPy.expand, &, series3 - series3_reference)

    # tested with the Python package BSeries
    series4 = modified_equation(f, u, dt, A, b, c, 4)
    series4_reference = [
      -dt^3*(-2*p^2*q*(2 - q)*(p - 1) + 2*p*q*(2 - q)*(p - 1)*(q - 2))/12 - dt^3*(-p^2*q*(2 - q)^2 + p*q^2*(p - 1)^2 + p*q*(1 - p)*(2 - q)*(p - 1) + p*q*(2 - q)*(p - 1)*(q - 2))/12 - dt^3*(p^2*q^2*(p - 1) - 2*p^2*q*(2 - q)^2 - p^2*q*(2 - q)*(p - 1) - p*q*(2 - q)^2*(p - 1) - p*q*(2 - q)*(p - 1)^2 - p*q*(p - 1)^3 + p*(2 - q)^4)/4 - dt^2*p*q*(2 - q)*(p - 1)/6 + dt^2*(-p^2*q*(2 - q) - p*q*(2 - q)*(p - 1) - p*q*(p - 1)^2 + p*(2 - q)^3)/3 - dt*(-p*q*(p - 1) + p*(2 - q)^2)/2 + p*(2 - q),
        -dt^3*(-2*p*q^2*(2 - q)*(p - 1) + 2*p*q*(2 - q)*(p - 1)^2)/12 - dt^3*(p^2*q*(2 - q)^2 - p*q^2*(p - 1)^2 + p*q*(2 - q)^2*(p - 1) + p*q*(2 - q)*(p - 1)^2)/12 - dt^3*(-p^2*q^2*(2 - q) - p*q^2*(2 - q)*(p - 1) - 2*p*q^2*(p - 1)^2 + p*q*(2 - q)^3 + p*q*(2 - q)^2*(p - 1) + p*q*(2 - q)*(p - 1)^2 + q*(p - 1)^4)/4 + dt^2*p*q*(2 - q)*(p - 1)/6 + dt^2*(-p*q^2*(p - 1) + p*q*(2 - q)^2 + p*q*(2 - q)*(p - 1) + q*(p - 1)^3)/3 - dt*(p*q*(2 - q) + q*(p - 1)^2)/2 + q*(p - 1)
    ]
    @test mapreduce(iszero ∘ SymPy.expand, &, series4 - series4_reference)
  end

  @testset "Symbolics.jl" begin
    # Lotka-Volterra model
    Symbolics.@variables dt
    u = Symbolics.@variables p q
    f = [p * (2 - q), q * (p - 1)]

    # Explicit Euler method
    A = @SArray [0//1;]
    b = @SArray [1//1]
    c = @SArray [0//1]

    # tested with the Python package BSeries
    series2 = modified_equation(f, u, dt, A, b, c, 2)
    series2_reference = [
      -dt*(-p*q*(p - 1) + p*(2 - q)^2)/2 + p*(2 - q),
      -dt*( p*q*(2 - q) + q*(p - 1)^2)/2 + q*(p - 1)
    ]
    @test mapreduce(isequal, &, series2, series2_reference)

    # tested with the Python package BSeries
    series3 = modified_equation(f, u, dt, A, b, c, 3)
    series3_reference = [
      -dt^2*p*q*(2 - q)*(p - 1)/6 + dt^2*(-p^2*q*(2 - q) - p*q*(2 - q)*(p - 1) - p*q*(p - 1)^2 + p*(2 - q)^3)/3 - dt*(-p*q*(p - 1) + p*(2 - q)^2)/2 + p*(2 - q),
      dt^2*p*q*(2 - q)*(p - 1)/6 + dt^2*(-p*q^2*(p - 1) + p*q*(2 - q)^2 + p*q*(2 - q)*(p - 1) + q*(p - 1)^3)/3 - dt*(p*q*(2 - q) + q*(p - 1)^2)/2 + q*(p - 1)
    ]
    @test mapreduce(iszero ∘ Symbolics.expand, &, series3 - series3_reference)

    # tested with the Python package BSeries
    series4 = modified_equation(f, u, dt, A, b, c, 4)
    series4_reference = [
      -dt^3*(-2*p^2*q*(2 - q)*(p - 1) + 2*p*q*(2 - q)*(p - 1)*(q - 2))/12 - dt^3*(-p^2*q*(2 - q)^2 + p*q^2*(p - 1)^2 + p*q*(1 - p)*(2 - q)*(p - 1) + p*q*(2 - q)*(p - 1)*(q - 2))/12 - dt^3*(p^2*q^2*(p - 1) - 2*p^2*q*(2 - q)^2 - p^2*q*(2 - q)*(p - 1) - p*q*(2 - q)^2*(p - 1) - p*q*(2 - q)*(p - 1)^2 - p*q*(p - 1)^3 + p*(2 - q)^4)/4 - dt^2*p*q*(2 - q)*(p - 1)/6 + dt^2*(-p^2*q*(2 - q) - p*q*(2 - q)*(p - 1) - p*q*(p - 1)^2 + p*(2 - q)^3)/3 - dt*(-p*q*(p - 1) + p*(2 - q)^2)/2 + p*(2 - q),
        -dt^3*(-2*p*q^2*(2 - q)*(p - 1) + 2*p*q*(2 - q)*(p - 1)^2)/12 - dt^3*(p^2*q*(2 - q)^2 - p*q^2*(p - 1)^2 + p*q*(2 - q)^2*(p - 1) + p*q*(2 - q)*(p - 1)^2)/12 - dt^3*(-p^2*q^2*(2 - q) - p*q^2*(2 - q)*(p - 1) - 2*p*q^2*(p - 1)^2 + p*q*(2 - q)^3 + p*q*(2 - q)^2*(p - 1) + p*q*(2 - q)*(p - 1)^2 + q*(p - 1)^4)/4 + dt^2*p*q*(2 - q)*(p - 1)/6 + dt^2*(-p*q^2*(p - 1) + p*q*(2 - q)^2 + p*q*(2 - q)*(p - 1) + q*(p - 1)^3)/3 - dt*(p*q*(2 - q) + q*(p - 1)^2)/2 + q*(p - 1)
    ]
    @test mapreduce(iszero ∘ Symbolics.expand, &, series4 - series4_reference)
  end
end

@testset "modifying_integrator" begin
  t1  = rootedtree([1])
  t2  = rootedtree([1, 2])
  t31 = rootedtree([1, 2, 2])
  t32 = rootedtree([1, 2, 3])
  t41 = rootedtree([1, 2, 2, 2])
  t42 = rootedtree([1, 2, 2, 3])
  t43 = rootedtree([1, 2, 3, 3])
  t44 = rootedtree([1, 2, 3, 4])
  t51 = rootedtree([1, 2, 2, 2, 2])
  t52 = rootedtree([1, 2, 2, 2, 3])
  t53 = rootedtree([1, 2, 2, 3, 3])
  t54 = rootedtree([1, 2, 2, 3, 4])
  t55 = rootedtree([1, 2, 3, 2, 3])
  t56 = rootedtree([1, 2, 3, 3, 3])
  t57 = rootedtree([1, 2, 3, 3, 4])
  t58 = rootedtree([1, 2, 3, 4, 4])
  t59 = rootedtree([1, 2, 3, 4, 5])

  @testset "SSPRK(2, 2) aka Heun's method" begin
    A = @SArray [0.0 0.0;
                1.0 0.0]
    b = @SArray [1/2, 1/2]
    c = @SArray [0.0, 1.0]
    order = 5
    series = modifying_integrator(A, b, c, order)

    # tested with the Python package BSeries
    @test series[t1 ] ≈ 1                    atol=10*eps()
    @test series[t2 ] ≈ 0                    atol=10*eps()
    @test series[t31] ≈ -0.166666666666667   atol=10*eps()
    @test series[t32] ≈ 1/6                  atol=10*eps()
    @test series[t41] ≈ 8.32667268468867e-17 atol=10*eps()
    @test series[t42] ≈ 0.125000000000000    atol=10*eps()
    @test series[t43] ≈ 1.38777878078145e-17 atol=10*eps()
    @test series[t44] ≈ -0.125000000000000   atol=10*eps()
    @test series[t51] ≈ 0.200000000000000    atol=10*eps()
    @test series[t52] ≈ 0.0583333333333333   atol=10*eps()
    @test series[t55] ≈ 0.00833333333333335  atol=10*eps()
    @test series[t53] ≈ -0.0583333333333333  atol=10*eps()
    @test series[t56] ≈ -0.200000000000000   atol=10*eps()
    @test series[t54] ≈ -0.133333333333333   atol=10*eps()
    @test series[t57] ≈ -0.0583333333333333  atol=10*eps()
    @test series[t58] ≈ 0.0583333333333333   atol=10*eps()
    @test series[t59] ≈ 0.133333333333333    atol=10*eps()
  end

  @testset "SSPRK(3, 3)" begin
    A = @SArray [0.0 0.0 0.0;
                1.0 0.0 0.0;
                1/4 1/4 0.0]
    b = @SArray [1/6, 1/6, 2/3]
    c = @SArray [0.0, 1.0, 1/2]
    order = 5
    series = modifying_integrator(A, b, c, order)

    # tested with the Python package BSeries
    @test series[t1 ] ≈ 1                    atol=10*eps()
    @test series[t2 ] ≈ 0                    atol=10*eps()
    @test series[t31] ≈ 0                    atol=10*eps()
    @test series[t32] ≈ 0                    atol=10*eps()
    @test series[t41] ≈ 0                    atol=10*eps()
    @test series[t42] ≈ 0.0416666666666667   atol=10*eps()
    @test series[t43] ≈ -0.0833333333333333  atol=10*eps()
    @test series[t44] ≈ 1/24                 atol=10*eps()
    @test series[t51] ≈ -0.00833333333333330 atol=10*eps()
    @test series[t52] ≈ 0.0166666666666667   atol=10*eps()
    @test series[t55] ≈ -0.0333333333333333  atol=10*eps()
    @test series[t53] ≈ -0.0166666666666667  atol=10*eps()
    @test series[t56] ≈ 0.00833333333333332  atol=10*eps()
    @test series[t54] ≈ -0.00833333333333334 atol=10*eps()
    @test series[t57] ≈ 0.0250000000000000   atol=10*eps()
    @test series[t58] ≈ 1/60                 atol=10*eps()
    @test series[t59] ≈ -0.0333333333333333  atol=10*eps()
  end

  @testset "Explicit midpoint method (Runge's method)" begin
    A = @SArray [0 0; 1//2 0]
    b = @SArray [0, 1//1]
    c = @SArray [0, 1//2]
    order = 5
    series = modifying_integrator(A, b, c, order)

    # tested with the Python package BSeries
    @test series[t1 ] == 1
    @test series[t2 ] == 0
    @test series[t31] == 1//12
    @test series[t32] == 1//6
    @test series[t41] == 0
    @test series[t42] == 0
    @test series[t43] == -1//8
    @test series[t44] == -1//8
    @test series[t51] == 1//80
    @test series[t52] == 1//60
    @test series[t55] == 7//240
    @test series[t53] == 1//240
    @test series[t56] == 9//80
    @test series[t54] == 1//80
    @test series[t57] == 13//120
    @test series[t58] == 13//80
    @test series[t59] == 2//15
  end
end

@testset "modifying_integrator with elementary differentials" begin
  @testset "SymEngine.jl" begin
    dt = SymEngine.symbols("dt")
    α, β, γ = SymEngine.symbols("alpha, beta, gamma")
    u1, u2, u3 = u = SymEngine.symbols("u1 u2 u3")
    f = [α*u[2]*u[3], β*u[3]*u[1], γ*u[1]*u[2]]

    # Implicit midpoint method
    A = @SArray [1//2;]
    b = @SArray [1//1]
    c = @SArray [1//2]

    # See eq. (12) of
    # - Philippe Chartier, Ernst Hairer and Gilles Vilmart (2007)
    #   Numerical integrators based on modified differential equations
    #   [DOI: 10.1090/S0025-5718-07-01967-9](https://doi.org/10.1090/S0025-5718-07-01967-9)

    series = modifying_integrator(f, u, dt, A, b, c, 5)
    # Dirty workaround since there is no way to get the polynomial coefficients
    # at the moment.
    # We differentiate the expression and set `dt` to zero to get the corresponding
    # coefficient divided by factorial(how often we needed to differentiate).
    s3_reference = -(α * β * u3^2 + α * γ * u2^2 + β * γ * u1^2) / 12
    for i in eachindex(f)
      s3 = SymEngine.subs(
        1//2 * SymEngine.diff(SymEngine.diff((series[i] - f[i]) / f[i], dt), dt),
        dt => 0)
      @test iszero(SymEngine.expand(s3 - s3_reference))
    end

    s5_reference = 6//5 * s3_reference^2 + 1//60 * α * β * γ * (
      β * u1^2 * u3^2 + γ * u2^2 * u1^2 + α * u3^2 * u2^2
    )
    for i in eachindex(f)
      s5 = SymEngine.subs(
        1//24 * SymEngine.diff(SymEngine.diff(SymEngine.diff(SymEngine.diff(
          (series[i] - f[i]) / f[i], dt), dt), dt), dt),
        dt => 0)
      @test iszero(SymEngine.expand(s5 - s5_reference))
    end
  end

  @testset "SymPy.jl" begin
    dt = SymPy.symbols("dt")
    α, β, γ = SymPy.symbols("alpha, beta, gamma")
    u1, u2, u3 = u = SymPy.symbols("u1 u2 u3")
    f = [α*u[2]*u[3], β*u[3]*u[1], γ*u[1]*u[2]]

    # Implicit midpoint method
    A = @SArray [1//2;]
    b = @SArray [1//1]
    c = @SArray [1//2]

    # See eq. (12) of
    # - Philippe Chartier, Ernst Hairer and Gilles Vilmart (2007)
    #   Numerical integrators based on modified differential equations
    #   [DOI: 10.1090/S0025-5718-07-01967-9](https://doi.org/10.1090/S0025-5718-07-01967-9)

    series = modifying_integrator(f, u, dt, A, b, c, 5)
    # Dirty workaround used also for the other symbolic setups - just make
    # it consistent here, although we could use another approach with SymPy.jl.
    # We differentiate the expression and set `dt` to zero to get the corresponding
    # coefficient divided by factorial(how often we needed to differentiate).
    s3_reference = -(α * β * u3^2 + α * γ * u2^2 + β * γ * u1^2) / 12
    for i in eachindex(f)
      s3 = SymPy.subs(
        1//2 * SymPy.diff(SymPy.diff((series[i] - f[i]) / f[i], dt), dt),
        dt => 0)
      @test iszero(SymPy.expand(s3 - s3_reference))
    end

    s5_reference = 6//5 * s3_reference^2 + 1//60 * α * β * γ * (
      β * u1^2 * u3^2 + γ * u2^2 * u1^2 + α * u3^2 * u2^2
    )
    for i in eachindex(f)
      s5 = SymPy.subs(
        1//24 * SymPy.diff(SymPy.diff(SymPy.diff(SymPy.diff(
          (series[i] - f[i]) / f[i], dt), dt), dt), dt),
        dt => 0)
      @test iszero(SymPy.expand(s5 - s5_reference))
    end
  end

  @testset "Symbolics.jl" begin
    Symbolics.@variables dt
    Symbolics.@variables α β γ
    u = Symbolics.@variables u1 u2 u3
    f = [α*u[2]*u[3], β*u[3]*u[1], γ*u[1]*u[2]]

    # Implicit midpoint method
    A = @SArray [1//2;]
    b = @SArray [1//1]
    c = @SArray [1//2]

    # See eq. (12) of
    # - Philippe Chartier, Ernst Hairer and Gilles Vilmart (2007)
    #   Numerical integrators based on modified differential equations
    #   [DOI: 10.1090/S0025-5718-07-01967-9](https://doi.org/10.1090/S0025-5718-07-01967-9)

    series = modifying_integrator(f, u, dt, A, b, c, 5)
    # Dirty workaround since there is no way to get the polynomial coefficients
    # in Symbolics.jl at the moment, see
    # https://github.com/JuliaSymbolics/Symbolics.jl/issues/216
    # We differentiate the expression and set `dt` to zero to get the corresponding
    # coefficient divided by factorial(how often we needed to differentiate).
    d = Symbolics.Differential(dt)
    s3_reference = -(α * β * u3^2 + α * γ * u2^2 + β * γ * u1^2) / 12
    for i in eachindex(f)
      s3 = Symbolics.simplify_fractions(Symbolics.substitute(
        Symbolics.expand_derivatives(1//2 * d(d((series[i] - f[i]) / f[i]))),
        dt => 0))
      @test isequal(s3, s3_reference)
    end

    s5_reference = 6//5 * s3_reference^2 + 1//60 * α * β * γ * (
      β * u1^2 * u3^2 + γ * u2^2 * u1^2 + α * u3^2 * u2^2
    )
    for i in eachindex(f)
      s5 = Symbolics.simplify_fractions(Symbolics.substitute(
        Symbolics.expand_derivatives(1//24 * d(d(d(d((series[i] - f[i]) / f[i]))))),
        dt => 0))
      @test iszero(Symbolics.expand(s5 - s5_reference))
    end
  end
end

@testset "nonlinear oscillator" begin
  @testset "SymEngine.jl" begin
    dt = SymEngine.symbols("dt")
    u = SymEngine.symbols("u_1, u_2")
    f = [-u[2], u[1]] / (u[1]^2 + u[2]^2)

    # explicit midpoint method
    A = @SArray [0 0; 1//2 0]
    b = @SArray [0, 1//1]
    c = @SArray [0, 1//2];

    # tested with Mathematica using
    # ```
    # ClearAll["Global`*"];
    #
    # solution[t_,
    #   u0_] := {u0[[1]]*Cos[t/(u0[[1]]^2 + u0[[2]]^2)] -
    #   u0[[2]]*Sin[t/(u0[[1]]^2 + u0[[2]]^2)],
    #   u0[[2]]*Cos[t/(u0[[1]]^2 + u0[[2]]^2)] +
    #   u0[[1]]*Sin[t/(u0[[1]]^2 + u0[[2]]^2)]}
    #
    # factorF2[u_] := 1/(12*(u[[1]]^2 + u[[2]]^2)^2)
    # (*factorF2[u_]:=ff2*)
    # factorF4[u_] := 1/(20*(u[[1]]^2 + u[[2]]^2)^4)
    # (*factorF4[u_]:=ff4*)
    # factorF6[u_] := 127/(2016*(u[[1]]^2 + u[[2]]^2)^6)
    # (*factorF6[u_]:=ff6*)
    # factorF8[u_] := 1513/(12960*(u[[1]]^2 + u[[2]]^2)^8)
    # (*factorF8[u_]:=ff8*)
    #
    # factorU3[u_] := 0
    # (*factorU3[u_]:=fu3*)
    # factorU5[u_] := 1/(48*(u[[1]]^2 + u[[2]]^2)^6)
    # (*factorU5[u_]:=fu5*)
    # factorU7[u_] := 31/(640*(u[[1]]^2 + u[[2]]^2)^8)
    # (*factorU7[u_]:=fu7*)
    # factorU9[u_] := 8969/(80640*(u[[1]]^2 + u[[2]]^2)^10)
    # (*factorU9[u_]:=fu9*)
    #
    # newModifiedF[
    #   u_] :=
    # (1 + factorF2[u]*h^2 + factorF4[u]*h^4 + factorF6[u]*h^6)*
    #   originalF[
    #     u] + (factorU3[u]*h^3 + factorU5[u]*h^5 + factorU7[u]*h^7)*u
    #
    # u0 = {u01, u02};
    # f = newModifiedF;
    # y2 = u0 + h/2*f[u0];
    # unew = u0 + h*f[y2];
    # difference = Simplify@Series[unew - solution[h, u0], {h, 0, 10}];
    #
    # Simplify[difference[[1]]]
    # Simplify[difference[[2]]]
    # ```
    series = modifying_integrator(f, u, dt, A, b, c, 10)
    terms = SymEngine.subs.(series, (Dict(u[1] => 1//1, u[2] => 0//1), ))

    @test isequal(terms[1], 1//48 * dt^5 + 31//640 * dt^7 + 8969//80640 * dt^9)
    @test isequal(terms[2], 1 + 1//12 * dt^2 + 1//20 * dt^4 + 127//2016 * dt^6 + 1513//12960 * dt^8)
  end
end

@testset "integer coefficients" begin
  # reported by David Ketcheson on 2021-12-09
  A = reshape([1], 1, 1); b = [1]; c = [1]
  @inferred modified_equation(A, b, c, 4)
end

@testset "Runge-Kutta methods interface" begin
  # classical RK4
  A = [0 0 0 0
       1//2 0 0 0
       0 1//2 0 0
       0 0 1 0]
  b = [1//6, 1//3, 1//3, 1//6]
  rk = RungeKuttaMethod(A, b)

  # fourth-order accurate
  series_integrator = @inferred bseries(rk, 4)
  series_exact = @inferred ExactSolution(series_integrator)
  @test mapreduce(==, &, values(series_integrator), values(series_exact))

  # not fifth-order accurate
  series_integrator = @inferred bseries(rk, 5)
  series_exact = @inferred ExactSolution(series_integrator)
  @test mapreduce(==, &, values(series_integrator), values(series_exact)) == false
end

@testset "additive Runge-Kutta methods interface" begin
  @testset "IMEX Euler (partitioned Euler, symplectic Euler)" begin
    ex_euler = RungeKuttaMethod(
      @SMatrix([0]), @SVector [1]
    )
    im_euler = RungeKuttaMethod(
      @SMatrix([1]), @SVector [1]
    )
    ark = AdditiveRungeKuttaMethod([im_euler, ex_euler])

    # first-order accurate
    series_integrator = @inferred bseries(ark, 1)
    series_exact = @inferred ExactSolution(series_integrator)
    @test mapreduce(==, &, values(series_integrator), values(series_exact))

    # not second-order accurate
    series_integrator = @inferred bseries(ark, 2)
    series_exact = @inferred ExactSolution(series_integrator)
    @test mapreduce(==, &, values(series_integrator), values(series_exact)) == false

    # modified equations and modifying integrators
    let series_integrator = @inferred bseries(ark, 3)
      @testset "modified_equation" begin
        mod_eq = @inferred modified_equation(series_integrator)

        # Hairer, Lubich, Wanner (2006) Geometric numerical integration
        # Table IX.10.1, p. 383
        # Black nodes are `0`, white nodes are `1`
        mod_eq_reference = Dict(
          rootedtree(Int[], Bool[])            => 0//1,
          rootedtree([1], Bool[0])             => 1//1,
          rootedtree([1], Bool[1])             => 1//1,
          rootedtree([1, 2], Bool[0, 0])       => 1//2,
          rootedtree([1, 2], Bool[1, 0])       => 1//2,
          rootedtree([1, 2], Bool[0, 1])       => -1//2,
          rootedtree([1, 2], Bool[1, 1])       => -1//2,
          rootedtree([1, 2, 3], Bool[0, 0, 0]) => 1//3,
          rootedtree([1, 2, 3], Bool[1, 0, 0]) => 1//3,
          rootedtree([1, 2, 3], Bool[0, 1, 0]) => -1//6,
          rootedtree([1, 2, 3], Bool[1, 1, 0]) => -1//6,
          rootedtree([1, 2, 3], Bool[0, 0, 1]) => -1//6,
          rootedtree([1, 2, 3], Bool[1, 0, 1]) => -1//6,
          rootedtree([1, 2, 3], Bool[0, 1, 1]) => 1//3,
          rootedtree([1, 2, 3], Bool[1, 1, 1]) => 1//3,
          rootedtree([1, 2, 2], Bool[0, 0, 0]) => 1//6,
          rootedtree([1, 2, 2], Bool[1, 0, 0]) => 1//6,
          rootedtree([1, 2, 2], Bool[0, 1, 0]) => -1//3,
          rootedtree([1, 2, 2], Bool[1, 1, 0]) => -1//3,
          rootedtree([1, 2, 2], Bool[0, 1, 1]) => 1//6,
          rootedtree([1, 2, 2], Bool[1, 1, 1]) => 1//6,
        )
        for (key, val) in mod_eq
          @test val == mod_eq_reference[key]
        end
      end

      @testset "modifying_integrator" begin
        mod_int = @inferred modifying_integrator(series_integrator)
      end
    end
  end

  @testset "Störmer-Verlet" begin
    # Hairer, Lubich, Wanner (2002)
    # Geometric numerical integration
    # Table II.2.1
    As = [
      [0 0; 1//2 1//2],
      [1//2 0; 1//2 0]
    ]
    bs = [
      [1//2, 1//2],
      [1//2, 1//2]
    ]
    ark = AdditiveRungeKuttaMethod(As, bs)

    # second-order accurate
    series_integrator = @inferred bseries(ark, 2)
    series_exact = @inferred ExactSolution(series_integrator)
    @test mapreduce(==, &, values(series_integrator), values(series_exact))

    # not third-order accurate
    series_integrator = @inferred bseries(ark, 3)
    series_exact = @inferred ExactSolution(series_integrator)
    @test mapreduce(==, &, values(series_integrator), values(series_exact)) == false

    # modified equations and modifying integrators
    let series_integrator = @inferred bseries(ark, 4)
      @testset "modified_equation" begin
        mod_eq = @inferred modified_equation(series_integrator)

        # second-order accurate
        for t in BicoloredRootedTreeIterator(2)
          @test iszero(mod_eq[t])
        end

        # Hairer, Lubich, Wanner (2006) Geometric numerical integration
        # Table IX.10.2, p. 386
        # Black nodes are `0`, white nodes are `1`
        series_reference = Dict(
          rootedtree(Int[], Bool[])            => 1//1,
          rootedtree([1], Bool[0])             => 1//1,
          rootedtree([1], Bool[1])             => 1//1,
          rootedtree([1, 2], Bool[0, 0])       => 1//2,
          rootedtree([1, 2], Bool[1, 0])       => 1//2,
          rootedtree([1, 2], Bool[0, 1])       => 1//2,
          rootedtree([1, 2], Bool[1, 1])       => 1//2,
          rootedtree([1, 2, 3], Bool[0, 0, 0]) => 1//4,
          rootedtree([1, 2, 3], Bool[1, 0, 0]) => 1//4,
          rootedtree([1, 2, 3], Bool[0, 1, 0]) => 0//1,
          rootedtree([1, 2, 3], Bool[1, 1, 0]) => 0//1,
          rootedtree([1, 2, 3], Bool[0, 0, 1]) => 1//4,
          rootedtree([1, 2, 3], Bool[1, 0, 1]) => 1//4,
          rootedtree([1, 2, 3], Bool[0, 1, 1]) => 1//4,
          rootedtree([1, 2, 3], Bool[1, 1, 1]) => 1//4,
          rootedtree([1, 2, 2], Bool[0, 0, 0]) => 1//2,
          rootedtree([1, 2, 2], Bool[1, 0, 0]) => 1//2,
          rootedtree([1, 2, 2], Bool[0, 1, 0]) => 1//4,
          rootedtree([1, 2, 2], Bool[1, 1, 0]) => 1//4,
          rootedtree([1, 2, 2], Bool[0, 1, 1]) => 1//4,
          rootedtree([1, 2, 2], Bool[1, 1, 1]) => 1//4,
        )
        for t in keys(series_reference)
          @test series_reference[t] == series_integrator[t]
        end
        mod_eq_reference = Dict(
          rootedtree(Int[], Bool[])            => 0//1,
          rootedtree([1], Bool[0])             => 1//1,
          rootedtree([1], Bool[1])             => 1//1,
          rootedtree([1, 2], Bool[0, 0])       => 0//1,
          rootedtree([1, 2], Bool[1, 0])       => 0//1,
          rootedtree([1, 2], Bool[0, 1])       => 0//1,
          rootedtree([1, 2], Bool[1, 1])       => 0//1,
          rootedtree([1, 2, 3], Bool[0, 0, 0]) => 1//12,
          rootedtree([1, 2, 3], Bool[1, 0, 0]) => 1//12,
          rootedtree([1, 2, 3], Bool[0, 1, 0]) => -1//6,
          rootedtree([1, 2, 3], Bool[1, 1, 0]) => -1//6,
          rootedtree([1, 2, 3], Bool[0, 0, 1]) => 1//12,
          rootedtree([1, 2, 3], Bool[1, 0, 1]) => 1//12,
          rootedtree([1, 2, 3], Bool[0, 1, 1]) => 1//12,
          rootedtree([1, 2, 3], Bool[1, 1, 1]) => 1//12,
          rootedtree([1, 2, 2], Bool[0, 0, 0]) => 1//6,
          rootedtree([1, 2, 2], Bool[1, 0, 0]) => 1//6,
          rootedtree([1, 2, 2], Bool[0, 1, 0]) => -1//12,
          rootedtree([1, 2, 2], Bool[1, 1, 0]) => -1//12,
          rootedtree([1, 2, 2], Bool[0, 1, 1]) => -1//12,
          rootedtree([1, 2, 2], Bool[1, 1, 1]) => -1//12,
        )
        for t in keys(mod_eq_reference)
          @test mod_eq_reference[t] == mod_eq[t]
        end

        # Hairer, Lubich, Wanner (2003)
        # Geometric numerical integration illustrated by the Störmer-Verlet method
        # https://doi.org/10.1017/S0962492902000144
        # equation (4.8)
        @testset "Pendulum, SymEngine" begin
          dt = SymEngine.symbols("dt")
          q, v = u = SymEngine.symbols("q, v")
          f = ([v, 0], [0, -sin(q)])

          series = modified_equation(f, u, dt, series_integrator)
          series_reference = f[1] + f[2] + 1//12 * dt^2 * [2 * cos(q) * v, sin(q) * cos(q) + sin(q) * v^2]
          @test mapreduce(iszero ∘ SymEngine.expand, &, series - series_reference)
        end
      end # @testset "modified_equation"

      @testset "modifying_integrator" begin
        mod_int = @inferred modifying_integrator(series_integrator)

        for t in BicoloredRootedTreeIterator(2)
          @test iszero(mod_int[t])
        end
      end
    end
  end

  @testset "SSPRK33 two times" begin
    # Using the same method multiple times is equivalent to using a plain RK
    # method without any splitting/partitioning/decomposition
    A = @SArray [0 0 0; 1 0 0; 1//4 1//4 0]
    b = @SArray [1//6, 1//6, 2//3]
    rk = RungeKuttaMethod(A, b)
    ark = AdditiveRungeKuttaMethod([rk, rk])

    # third-order accurate
    series_integrator = @inferred bseries(ark, 3)
    series_exact = @inferred ExactSolution(series_integrator)
    @test mapreduce(==, &, values(series_integrator), values(series_exact))

    # not fourth-order accurate
    series_integrator = @inferred bseries(ark, 4)
    series_exact = @inferred ExactSolution(series_integrator)
    @test mapreduce(==, &, values(series_integrator), values(series_exact)) == false

    # modified equations and modifying integrators
    let series_rk = bseries(rk, 5)
      series_ark = bseries(ark, order(series_rk))

      for colored_tree in keys(series_ark)
        tree = rootedtree(colored_tree.level_sequence)
        @test series_rk[tree] == series_ark[colored_tree]
      end

      @testset "modified_equation" begin
        mod_eq_rk  = modified_equation(series_rk)
        mod_eq_ark = modified_equation(series_ark)
        for colored_tree in keys(mod_eq_ark)
          tree = rootedtree(colored_tree.level_sequence)
          @test mod_eq_rk[tree] == mod_eq_ark[colored_tree]
        end
      end

      @testset "modifying_integrator" begin
        mod_int_rk  = modifying_integrator(series_rk)
        mod_int_ark = modifying_integrator(series_ark)
        for colored_tree in keys(mod_int_ark)
          tree = rootedtree(colored_tree.level_sequence)
          @test mod_int_rk[tree] == mod_int_ark[colored_tree]
        end
      end
    end
  end
end

@testset "multirate infinitesimal split methods interface" begin
  # NOTE: This is still considered experimental and might change at any time!

  @testset "KW43" begin
    # Oswald Knoth, Ralf Wolke
    # Implicit-explicit Runge-Kutta methods for computing atmospheric reactive flows
    # Applied Numerical Mathematics, Volume 28, (1998) Issue 2-4
    # https://doi.org/10.1016/S0168-9274(98)00051-8
    ns = 4
    A = zeros(Polynomial{Rational{Int}, :x}, ns+1, ns)
    D = zeros(Int, ns+1, ns)
    G = zeros(Rational{Int}, ns+1, ns)
    c = zeros(Rational{Int}, ns+1)
    dts = zeros(Rational{Int}, ns+1)

    A[2, 1] = 1 // 2
    A[3, 1] = -1 // 6
    A[3, 2] = 2 // 3
    A[4, 1] = 1 // 3
    A[4, 2] = -1 // 3
    A[4, 3] = 1
    A[5, 1] = 1 // 6
    A[5, 2] = 1 // 3
    A[5, 3] = 1 // 3
    A[5, 4] = 1 // 6
    A[5, :] = A[5, :] - A[4, :]
    A[4, :] = A[4, :] - A[3, :]
    A[3, :] = A[3, :] - A[2, :]

    c[2] = 1 // 2
    c[3] = 1 // 2
    c[4] = 1
    c[5] = 1

    D[2,1] = 1
    D[3,2] = 1
    D[4,3] = 1
    D[5,4] = 1

    dts[1] = 0
    dts[2] = 1 // 2
    dts[3] = 0
    dts[4] = 1 // 2
    dts[5] = 0

    mis = @inferred MultirateInfinitesimalSplitMethod(A, D, G, c)

    # third-order accurate
    series_integrator = @inferred bseries(mis, 3)
    series_exact = @inferred ExactSolution(series_integrator)
    @test mapreduce(==, &, values(series_integrator), values(series_exact))

    # not fourth-order accurate
    series_integrator = @inferred bseries(mis, 4)
    series_exact = @inferred ExactSolution(series_integrator)
    @test mapreduce(==, &, values(series_integrator), values(series_exact)) == false
  end

  @testset "MRI-GARK-ERK33" begin
    # Sandu, Adrian.
    # "A class of multirate infinitesimal GARK methods".
    # SIAM Journal on Numerical Analysis 57, no. 5 (2019): 2300-2327.
    # https://doi.org/10.1137/18M1205492
    ns = 3
    A = zeros(Polynomial{Float64, :x}, ns+1, ns)
    D = zeros(ns+1, ns)
    G = zeros(ns+1, ns)
    c = zeros(ns+1)
    dts = zeros(ns+1)

    delta = -0.5

    # Γ^i in Sandu (2019) is the i-th term
    A[2, 1] = Polynomial([1 / 3])
    A[3, 1] = Polynomial([(-6 * delta - 7) / 12, (2 * delta + 1) / 2])
    A[3, 2] = Polynomial([(6 * delta + 11) / 12, -(2 * delta + 1) / 2])
    A[4, 1] = Polynomial([0.0, 0.5])
    A[4, 2] = Polynomial([(6 * delta - 5) / 12, -(2 * delta + 1) / 2])
    A[4, 3] = Polynomial([(3 - 2 * delta) / 4, delta])

    # These should be the usual `c` coefficients → c[3] = 2 / 3?
    c[2] = 1 / 3
    c[3] = 2 / 3
    c[4] = 1

    # All one???
    D[2, 1] = 1
    D[3, 2] = 1
    D[4, 3] = 1

    # Differences of `c`s?
    dts[1] = 0
    dts[2] = 1 / 3
    dts[3] = 1 / 3
    dts[4] = 1 / 3

    mis = @inferred MultirateInfinitesimalSplitMethod(A, D, G, c)

    # third-order accurate
    series_integrator = @inferred bseries(mis, 3)
    series_exact = @inferred ExactSolution(series_integrator)
    @test mapreduce(≈, &, values(series_integrator), values(series_exact))

    # not fourth-order accurate
    series_integrator = @inferred bseries(mis, 4)
    series_exact = @inferred ExactSolution(series_integrator)
    @test mapreduce(≈, &, values(series_integrator), values(series_exact)) == false
  end

  @testset "MRI-GARK-ERK45a" begin
    # Sandu, Adrian.
    # "A class of multirate infinitesimal GARK methods".
    # SIAM Journal on Numerical Analysis 57, no. 5 (2019): 2300-2327.
    # https://doi.org/10.1137/18M1205492
    ns = 5
    A = zeros(Polynomial{Float64, :x}, ns+1, ns)
    D = zeros(ns+1, ns)
    G = zeros(ns+1, ns)
    c = zeros(ns+1)
    dts = zeros(ns+1)

    # Γ^i in Sandu (2019) is the i-th term
    A[2,1] = Polynomial([1 / 5])

    A[3,1] = Polynomial([-53 / 16, 503 / 80])
    A[3,2] = Polynomial([281 / 80, -503 / 80])

    A[4,1] = Polynomial([-36562993.0/71394880.0, -1365537.0/35697440.0])
    A[4,2] = Polynomial([34903117.0/17848720.0, 4963773.0/7139488.0])
    A[4,3] = Polynomial([-88770499.0/71394880.0, -1465833.0/2231090.0])

    A[5,1] = Polynomial([-7631593.0/71394880.0, 66974357.0/35697440.0])
    A[5,2] = Polynomial([-166232021.0/35697440.0, 21445367.0/7139488.0])
    A[5,3] = Polynomial([6068517.0/1519040.0, -3.0])
    A[5,4] = Polynomial([8644289.0/8924360.0, -8388609.0/4462180.0])

    A[6,1] = Polynomial([277061.0/303808.0, -18227.0/7520.0])
    A[6,2] = Polynomial([-209323.0/1139280.0, 2.0])
    A[6,3] = Polynomial([-1360217.0/1139280.0, 1.0])
    A[6,4] = Polynomial([-148789.0/56964.0, 5.0])
    A[6,5] = Polynomial([147889.0/45120.0, -41933.0/7520.0])

    # These should be the usual `c` coefficients
    c[2] = 1 / 5
    c[3] = 2 / 5
    c[4] = 3 / 5
    c[5] = 4 / 5
    c[6] = 1

    D[2,1] = 0
    D[3,2] = 1
    D[4,3] = 1
    D[5,4] = 1
    D[6,5] = 1

    # Differences of `c`s?
    dts[1] = 0.0
    dts[2] = 1 / 5
    dts[3] = 1 / 5
    dts[4] = 1 / 5
    dts[5] = 1 / 5
    dts[6] = 1 / 5

    mis = @inferred MultirateInfinitesimalSplitMethod(A, D, G, c)

    # fourth-order accurate
    series_integrator = @inferred bseries(mis, 4)
    series_exact = @inferred ExactSolution(series_integrator)
    @test mapreduce(≈, &, values(series_integrator), values(series_exact))

    # not fifth-order accurate
    series_integrator = @inferred bseries(mis, 5)
    series_exact = @inferred ExactSolution(series_integrator)
    @test mapreduce(≈, &, values(series_integrator), values(series_exact)) == false
  end

  @testset "EB4" begin
    # Krogstad, Stein
    # "Generalized integrating factor methods for stiff PDEs".
    # Journal of Computational Physics 203, no. 1 (2005): 72-88.
    # https://doi.org/10.1016/j.jcp.2004.08.006
    ns = 4
    A = zeros(Polynomial{Float64, :x}, ns+1, ns)
    D = zeros(ns+1, ns)
    G = zeros(ns+1, ns)
    c = zeros(ns+1)
    dts = zeros(ns+1)

    fac1 = 1.0
    fac2 = 2.0

    A[2,1] = Polynomial([0.5])

    A[3,1] = Polynomial([0.5, -1 / fac1])
    A[3,2] = Polynomial([0.0, 1 / fac1])

    A[4,1] = Polynomial([1.0, -2 / fac1])
    A[4,3] = Polynomial([0.0, 2 / fac1])

    A[5,1] = Polynomial([1.0, -3 / fac1, 4 / fac2])
    A[5,2] = Polynomial([0.0, 2 / fac1, -4 / fac2])
    A[5,3] = Polynomial([0.0, 2 / fac1, -4 / fac2])
    A[5,4] = Polynomial([0.0, -1 / fac1, 4 / fac2])

    c[1] = 0
    c[2] = 0.5
    c[3] = 0.5
    c[4] = 1
    c[5] = 1

    mis = @inferred MultirateInfinitesimalSplitMethod(A, D, G, c)

    # fourth-order accurate
    series_integrator = @inferred bseries(mis, 4)
    series_exact = @inferred ExactSolution(series_integrator)
    @test mapreduce(≈, &, values(series_integrator), values(series_exact))

    # not fifth-order accurate
    series_integrator = @inferred bseries(mis, 5)
    series_exact = @inferred ExactSolution(series_integrator)
    @test mapreduce(≈, &, values(series_integrator), values(series_exact)) == false
  end

  @testset "OwrenRK4" begin
    # Owren, Brynjulf.
    # "Order conditions for commutator-free Lie group methods".
    # Journal of Physics A: Mathematical and General 39, no. 19 (2006): 5585.
    # https://doi.org/10.1088/0305-4470/39/19/S15
    ns = 5
    A = zeros(Polynomial{Float64, :x}, ns+1, ns)
    D = zeros(ns+1, ns)
    G = zeros(ns+1, ns)
    c = zeros(ns+1)
    dts = zeros(ns+1)

    A[2,1] = Polynomial([1 / 2])

    A[3,2] = Polynomial([1 / 2])

    A[4,1] = Polynomial([-1 / 2])
    A[4,3] = Polynomial([1.0])

    A[5,1] = Polynomial([1 / 4])
    A[5,2] = Polynomial([1 / 6])
    A[5,3] = Polynomial([1 / 6])
    A[5,4] = Polynomial([-1 / 12])

    A[6,1] = Polynomial([-1 / 12])
    A[6,2] = Polynomial([1 / 6])
    A[6,3] = Polynomial([1 / 6])
    A[6,4] = Polynomial([1 / 4])

    D[4,2] = 1
    D[6,5] = 1

    dts[1] = 0
    dts[2] = 1 / 2
    dts[3] = 1 / 2
    dts[4] = 1 / 2
    dts[5] = 1 / 2
    dts[6] = 1 / 2

    mis = @inferred MultirateInfinitesimalSplitMethod(A, D, G, c)

    # fourth-order accurate
    series_integrator = @inferred bseries(mis, 4)
    series_exact = @inferred ExactSolution(series_integrator)
    @test mapreduce(≈, &, values(series_integrator), values(series_exact))

    # not fifth-order accurate
    series_integrator = @inferred bseries(mis, 5)
    series_exact = @inferred ExactSolution(series_integrator)
    @test mapreduce(≈, &, values(series_integrator), values(series_exact)) == false
  end

  @testset "MISBI4_4" begin
    # Knoth, Oswald, and Joerg Wensch.
    # "Generalized split-explicit Runge–Kutta methods for the compressible Euler equations".
    # Monthly Weather Review 142, no. 5 (2014): 2067-2081.
    # https://doi.org/10.1175/MWR-D-13-00068.1
    # Table 4, method MIS4a
    ns = 4
    A = zeros(Polynomial{Float64, :x}, ns+1, ns)
    D = zeros(ns+1, ns)
    G = zeros(ns+1, ns)
    c = zeros(ns+1)
    dts = zeros(ns+1)

    # β in Knoth, Wensch (2014)
    A[2,1] = Polynomial([1 / 2])
    A[2,1] = Polynomial([0.38758444641450318])
    A[3,1] = Polynomial([-2.5318448354142823e-002])
    A[3,2] = Polynomial([0.38668943087310403])
    A[4,1] = Polynomial([0.20899983523553325])
    A[4,2] = Polynomial([-0.45856648476371231])
    A[4,3] = Polynomial([0.43423187573425748])
    A[5,1] = Polynomial([-0.10048822195663100])
    A[5,2] = Polynomial([-0.46186171956333327])
    A[5,3] = Polynomial([0.83045062122462809])
    A[5,4] = Polynomial([0.27014914900250392])

    # seems to be unused?
    c[2] = 0.38758444641450318
    c[3] = 0.61521685655017821
    c[4] = 0.23254717315441453
    c[5] = 1.0000000000000002

    # α in Knoth, Wensch (2014)
    D[3,2] =  0.52349249922385610
    D[4,2] =  1.1683374366893629
    D[4,3] = -0.75762080241712637
    D[5,2] = -3.6477233846797109e-002
    D[5,3] =  0.56936148730740477
    D[5,4] =  0.47746263002599681

    # γ in Knoth, Wensch (2014)
    G[3,2] =  0.13145089796226542
    G[4,2] = -0.36855857648747881
    G[4,3] =  0.33159232636600550
    G[5,2] = -6.5767130537473045e-002
    G[5,3] =  4.0591093109036858e-002
    G[5,4] =  6.4902111640806712e-002

    for i in 1:ns+1
      for j in 1:i-1
        dts[i] = dts[i] + A[i,j][0]
      end
    end

    mis = @inferred MultirateInfinitesimalSplitMethod(A, D, G, c)

    # third-order accurate
    series_integrator = @inferred bseries(mis, 3)
    series_exact = @inferred ExactSolution(series_integrator)
    @test_broken mapreduce(≈, &, values(series_integrator), values(series_exact))

    # not fourth-order accurate
    series_integrator = @inferred bseries(mis, 4)
    series_exact = @inferred ExactSolution(series_integrator)
    @test mapreduce(≈, &, values(series_integrator), values(series_exact)) == false
  end
end

end # @testset "BSeries"


module MultiStochGrad

using Printf

using Logging
using Base.CoreLogging
using LinearAlgebra

debug_log = stdout
logger = ConsoleLogger(stdout, CoreLogging.Info)
global_logger(logger)

include("scsg.jl")
include("svrg.jl")

include("applis/LinearRegression.jl")
include("applis/LogisticRegression.jl")

# export list

export Observations,
    TermFunction,
    compute_value,
    TermGradient,
    compute_gradient!,
    Evaluator,
    get_nbterms,
    # from scsg
    SCSG,
    minimize,
    BatchSizeInfo,
    get_batchsizeinfo,
    #
    SVRG
end


using ParameterizedFunctions, DiffEqBase
using Test, LinearAlgebra, InteractiveUtils

using SpecialFunctions

### ODE Macros

println("Build some examples")
f_t = @ode_def SymCheck begin # Checks for error due to symbol on 1
  dx = x
  dy = -c*y + d*x*y*t^2
end a b c d

@test DiffEqBase.__has_syms(f_t)

f_t2 = @ode_def SymCheck2 begin # Checks for error due to symbol on 1
  dx = 1
  dy = -c*y + d*x*y*t^2
end a b c d

f_t3 = @ode_def ExprCheck begin # Checks for error due to symbol on 1
  dx = a*x - b*x*y
  dy = -c*y + d*x*y
end a b c d # Change to π after unicode fix

f = @ode_def_all LotkaVolterra begin
  dx = a*x - b*x*y
  dy = -c*y + d*x*y
end a b c d

f_2 = @ode_def LotkaVolterra3 begin
  dx = a*x - b^2*x*y
  dy = -c*y + d*x*y
end a b c d

println("Test Values")
t = 1.0
u = [2.0,3.0]
p = [1.5,1,3,1]
du = zeros(2)
grad = similar(du)
J = zeros(2,2)
iJ= zeros(2,2)
iW= zeros(2,2)
f(du,u,p,t)
@test du == [-3.0,-3.0]
@test du == f(u,p,t)
f_t(du,u,p,t)
@test du == [2.0,-3.0]
f_t2(du,u,p,t)
@test du == [1.0,-3.0]

println("Test t-gradient")
f.tgrad(grad,u,p,t)
@test grad == zeros(2)
f_t.tgrad(grad,u,p,t)
@test grad == [0.0;12.0]

println("Test Jacobians")
f.jac(J,u,p,t)
@test J  == [-1.5 -2.0
             3.0 -1.0]
#@test f.jac(u, p, t) == [-1.5 -2.0; 3.0 -1.0]

println("Test Inv Rosenbrock-W")
f.invW(iW,u,p,2.0,t)
@test minimum(iW - inv(I - 2*J) .< 1e-10)

f.invW_t(iW,u,p,2.0,t)
@test minimum(iW - inv(I/2 - J) .< 1e-10)

println("Parameter Jacobians")
pJ = zeros(Float64,2,4)
f.paramjac(pJ,u,[2.0;2.5;3.0;1.0],t)
@test pJ == [2.0 -6.0  0.0 0.0
             0.0  0.0 -3.0 6.0]

@code_llvm DiffEqBase.has_jac(f)

println("Test booleans")
@test DiffEqBase.has_jac(f) == true
@test DiffEqBase.has_paramjac(f) == true

@code_llvm DiffEqBase.has_paramjac(f)

println("Test difficult differentiable")
NJ = @ode_def DiffDiff begin
  dx = a*x - b*x*y
  dy = -c*y + erf(x*y/d)
end a b c d
NJ(du,u,[1.5,1,3,4],t)
@test du == [-3.0;-3*3.0 + erf(2.0*3.0/4)]
@test du == NJ(u, [1.5,1,3,4], t)
# NJ(Val{:jac},t,u,J) # Currently gives E not defined, will be fixed by the next SymEgine

test_fail(x,y,d) = erf(x*y/d)
println("Test non-differentiable")
NJ = @ode_def NoJacTest begin
  dx = a*x - μ*x*y
  dy = -c*y + test_fail(x,y,d)
end a μ c d
NJ(du,u,[1.5,1,3,4],t)
@test du == [-3.0;-3*3.0 + erf(2.0*3.0/4)]
@test du == NJ(u,[1.5,1,3,4],t)
@test_throws MethodError NJ.jac(iJ,u,p,t)
# NJ(Val{:jac},t,u,J) # Currently gives E not defined, will be fixed by the next SymEgine

println("Test anonymous definition")

f_t_noname = @ode_def begin # Checks for error due to symbol on 1
  dx = x
  dy = -c*y + d*x*y*t^2
end a b c d

@test DiffEqBase.__has_syms(f_t_noname)

f = @ode_def begin
  dx = a*x - b*x*y
  dy = -c*y + d*x*y
end a b c d
@test_nowarn f([0.1,0.2], [1,2], [1,2,3,4], 1)

# Test failures of derivatives should not have #undef
sir_ode = @ode_def SIRModel begin
  dS = -b*S*I
  dI =  b*S*I - g*I
  dR =  g*I
end b g


# ---
# title: 45. Jump Game II
# id: problem45
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Array, Greedy
# link: <https://leetcode.com/problems/jump-game-ii/description/>
# hidden: true
# ---
#
# Given an array of non-negative integers `nums`, you are initially positioned
# at the first index of the array.
#
# Each element in the array represents your maximum jump length at that
# position.
#
# Your goal is to reach the last index in the minimum number of jumps.
#
# You can assume that you can always reach the last index.
#
#
#
# **Example 1:**
#
#
#
#     Input: nums = [2,3,1,1,4]
#     Output: 2
#     Explanation: The minimum number of jumps to reach the last index is 2. Jump 1 step from index 0 to 1, then 3 steps to the last index.
#
#
# **Example 2:**
#
#
#
#     Input: nums = [2,3,0,1,4]
#     Output: 2
#
#
#
#
# **Constraints:**
#
#   * `1 <= nums.length <= 3 * 104`
#   * `0 <= nums[i] <= 105`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


#### Tolerances

export ResidualTolerance, SolutionTolerance

abstract type AbstractTolerance{FT <: AbstractFloat} end

"""
    ResidualTolerance

A tolerance type based on the residual of the equation ``f(x) = 0``
"""
struct ResidualTolerance{FT} <: AbstractTolerance{FT}
    tol::FT
end

"""
    (tol::ResidualTolerance)(x1, x2, y)

Evaluates residual tolerance, based on ``|f(x)|``
"""
(tol::ResidualTolerance)(x1, x2, y) = all(abs.(y) .< tol.tol)

"""
    SolutionTolerance

A tolerance type based on the solution ``x`` of the equation ``f(x) = 0``
"""
struct SolutionTolerance{FT} <: AbstractTolerance{FT}
    tol::FT
end

"""
    (tol::SolutionTolerance)(x1, x2, y)

Evaluates solution tolerance, based on ``|x2-x1|``
"""
(tol::SolutionTolerance)(x1, x2, y) = all(abs.(x2 .- x1) .< tol.tol)