using ApproxFunOrthogonalPolynomials, SpecialFunctions, LinearAlgebra, Test
import DomainSets
import DomainSets: UnionDomain
import ApproxFunBase: space, SpaceOperator,
                    testspace, testbandedoperator, testraggedbelowoperator, testcalculus, testtransforms,
                    testfunctional

@testset "Spaces" begin
    @testset "evaluate" begin
        @test evaluate([1],Legendre(),0)==evaluate([1.0],Legendre(),0)
    end
    @testset "ChebyshevDirichlet" begin
        testtransforms(ChebyshevDirichlet{1,1}())

        @test Fun(exp,ChebyshevDirichlet{1,1})(.1) ≈ exp(.1)
        @test Fun(Fun(exp,ChebyshevDirichlet{1,1}),Ultraspherical(1))(.1) ≈ exp(.1)

        @test domain(ChebyshevDirichlet{1,1}()) === ChebyshevInterval()
        @test Fun(Fun(exp,ChebyshevDirichlet{1,1}(Segment(-1,1))),Ultraspherical(1))(.1) ≈ exp(.1)

        d = ChebyshevInterval()
        sp = ChebyshevDirichlet{1,1}(d)
        B=Dirichlet(sp)
        D=Derivative(d)
        L=D^2+I
        @test B[1:2,1:4] ≈ [1 -1 0 0; 1 1 0 0]

        @test csc(2)sin(1 - 0.1)  ≈ ([Dirichlet(d);L]\[[1.,0.],0.])(0.1)
        @test csc(2)sin(1 - 0.1)  ≈ ([B;L]\[[1.,0.],0.])(0.1)

        @test norm(([B;L]\[[1.,0],0])-([Dirichlet(d);L]\[[1.,0],0])) <10eps()
    end

    @testset "PiecewiseSpace" begin
        x = Fun(identity,UnionDomain(-1..0, 0..1))
        sp=space(x)
        testtransforms(sp;minpoints=2)

        D = Derivative(sp)
        testbandedoperator(D)

        B = [Dirichlet(sp); continuity(sp,0:1)]
        u = [B; D^2] \ [[1,0],zeros(2),0];
        @time u2 = [Dirichlet();Derivative(Chebyshev())^2] \ [[1.,0],0]
        @test u(0.) ≈ u2(0.)

        x = Fun(identity, UnionDomain(-10..0, 0..1, 1..15))
        sp=space(x)
        D=Derivative(sp)
        B=Dirichlet(sp)

        @time  u=[B;
            continuity(sp,0:1);
            D^2-x]\[[airyai(-10.),0],zeros(4),0];

        @test u(0.) ≈ airyai(0.)

        s=Fun(sin,-2..2)|>abs
        c=Fun(cos,-2..2)|>abs
        sc=Fun(x -> abs(sin(x))+abs(cos(x)), ContinuousSpace(PiecewiseSegment([-2,(-π/2),0,π/2,2])))
        @test norm(sc-(c+s))<100eps()
    end

    @testset "max/min creates breakpoints" begin
        x = Fun()
        g = 4*(x-0.2)
        f = max(-1,g)
        f2 = min(f,1)

        @test norm(max(x,x)-x)<100eps()
        @test norm(min(x,x)-x)<100eps()

        f3=Fun(x->x<-0.05 ? -1.0 : (x<0.45 ? 4*(x-.2) : 1), (-1..1) \ [-0.05,0.45])
        @test norm(f2.(range(-1,stop=1,length=10))-f3.(range(-1,stop=1,length=10))) < 4eps()
    end

    @testset "Complex piecewise" begin
        x = Fun(identity, Segment(im,0) ∪ Segment(0,1))
        @test x(0.5) ≈ 0.5
        @test x(0.5im) ≈ 0.5im

        @test Fun(Fun(1.0),space(x))(0.5) == 1.0
        @test Fun(Fun(1.0),space(x))(0.5im) == 1.0

        @test (x+1)(0.5) ≈ 1.5
        @test (x-1)(0.5) ≈ -0.5
        @test (1-x)(0.5) ≈ 0.5
    end

    @testset "Equivalent spaces" begin
        @test norm(Fun(cos,Chebyshev)-Fun(cos,Jacobi(-0.5,-0.5)))<100eps()
    end

    @testset "ContinuousSpace" begin
        d = PiecewiseSegment(1.,2.,3.,4.)
        S = ContinuousSpace(d)
        testtransforms(S; minpoints=3, invertibletransform=false)

        D = Derivative(S)
        testbandedoperator(D)

        A = [ldirichlet(S); D-I]
        testraggedbelowoperator(A)
        @time u = [ldirichlet(S); D-I] \ [exp(1.); 0]

        @test u(1.1) ≈ exp(1.1)
        @test u(3.4) ≈ exp(3.4)
        @test last(u) ≈ exp(4)

        d=PiecewiseSegment(0,1.,1+im,im,0.)
        s=Space(d)

        # narrow down bug
        @test s isa ContinuousSpace
        @test ApproxFunBase.rangetype(s) == Float64
        cs=ApproxFunBase.canonicalspace(s)
        @test ApproxFunBase.rangetype(cs) == Float64

        @test conversion_type(s,cs) == s

        @test Fun(exp,d)(.1) ≈ exp(.1)
    end

    @testset "Jacobi recurrence bug" begin
        S=Jacobi(0.5,-0.5)
        f=Fun(exp,S)
        @test f(0.1) ≈ exp(0.1)
    end

    @testset "Jacobi evaluation bug" begin
        S=Jacobi(0.5,-0.5)
        f=Fun(exp,S)
        B = ldirichlet(S)
        testfunctional(B)
        @test B*f ≈ exp(-1)
    end

    @testset "Hermite" begin
        f=Fun(x->x+x^2,Hermite())
        @test f(1.) ≈ 2.
        @test values(f) ≈ points(f)+points(f).^2

        D = Derivative(Hermite())
        @time testbandedoperator(D)

        g = D*f
        @test g(1.) ≈ 3.
    end

    @testset "Arc exp" begin
        z=Fun(identity,Arc(0.,.1,0.,π/2))
        @test exp(z)(0.1exp(0.2im)) ≈ exp(0.1exp(0.2im))
    end

    @testset "Line" begin
        f=Fun(z->2exp(z^2),Line(0.,π/2))
        @test f(1.1im) ≈ 2exp(-1.1^2)
    end

    @testset "Exp Curve" begin
        a = 1+10*im; b = 2-6*im
        d = IntervalCurve(Fun(x->1+a*x+x^2+b*x^3))

        x=Fun(d)

        @test exp(x)(1+a*0.1+0.1^2+b*0.1^3) ≈ exp(1+a*0.1+0.1^2+b*0.1^3)
    end

    @testset "ChebyshevDirichlet multiplication" begin
        S=ChebyshevDirichlet()
        x=Fun()
        @test norm((ApproxFunOrthogonalPolynomials.Recurrence(S)*Fun(exp,S)-Fun(x->x*exp(x),S)).coefficients) < 100eps()
        @test norm((x*Fun(exp,S)-Fun(x->x*exp(x),S)).coefficients) < 100eps()

        @test Integral(S,1)*Fun(S,[1.,2.,3.]) ≈ integrate(Fun(Fun(S,[1.,2.,3.]),Chebyshev()))
    end

    @testset "QuotientSpace" begin
        S = Chebyshev()
        for B in (Dirichlet(S),
                    Neumann(S),
                    vcat(bvp(S,4)...),
                    [Dirichlet(S); Neumann(S); Evaluation(S, -1, 3); Evaluation(S, 1, 3)])
            QS = QuotientSpace(B)
            A = Conversion(QS, S)

            @test norm(B[1:size(B,1),1:10]*A[1:10,1:10-size(B,1)]) < 10000eps()
        end
    end

    @testset "Fast Ultraspherical*Chebyshev" begin
        f = Fun(exp, Ultraspherical(1))
        g = Fun(exp)

        @test ApproxFunBase.hasfasttransform(f)
        @test ApproxFunBase.pointscompatible(f,g)
        @test ApproxFunBase.hasfasttransformtimes(f,g)
        @test ApproxFunBase.transformtimes(f,g)(0.1) ≈ exp(0.2)
    end

    @testset "PointSpace" begin
        δ = Fun(ApproxFunBase.PointSpace([2.0]),[1.0])
        f = Fun(x->cos(50x)) + δ
        g = Fun(x->cos(50x),Ultraspherical(1)) + δ
        @test domain(f) == domain(g)

        @test (f*g)(0.1) ≈ cos(50*0.1)^2
        @test (f*g)(2.0) ≈ 1

        @test exp(f)(0.1) ≈ exp(cos(50*0.1))
        @test exp(f)(2.0) ≈ exp(1)

        @test iszero(sign(-Fun(zero)))
        @test iszero(sign(Fun(zero)))
        @test iszero(abs(-Fun(zero)))
        @test iszero(abs(Fun(zero)))
        @test angle(-Fun(zero)) ≈ Fun(π)
        @test iszero(angle(Fun(zero)))

        @test sign(f)(0.1) ≈ sign(cos(50*0.1))
        @test sign(f)(2.0) ≈ 1
        @test abs(f)(0.1) ≈ abs(cos(50*0.1))
        @test abs(f)(2.0) ≈ 1
        @test angle(f)(0.1) ≈ angle(cos(50*0.1))
        @test angle(f)(2.0) ≈ 0
    end

    @testset "Jump Locations" begin
        x = Fun(UnionDomain(0..1, 2..3))
        @test length(jumplocations(sign(x))) == 0

        x = Fun(Chebyshev(-1..1))
        @test length(jumplocations(x)) == 0
        @test all(jumplocations(sign(x) + sign(x+0.2)) .≈ [-0.2, 0])
    end

    @testset "SumSpace Conversion" begin
        H = ApproxFunBase.HeavisideSpace([-1.0,0.0,1.0])
        C = ContinuousSpace(ApproxFunBase.PiecewiseSegment([-1.0,0,1]))
        S = H + C
        P = Ultraspherical(1,Segment(-1.0,0.0)) ∪ Ultraspherical(1,Segment(0.0,1.0))
        f = Fun(S, randn(100))
        @test f(0.1) ≈ Fun(f, P)(0.1)

        @test Conversion(S,P)[1,1] == 1.0
        @test Conversion(S,P)[1,2] == 0.5
        @test rangespace(Conversion(S,P)) == Ultraspherical(1,Segment(-1.0,0.0)) ∪ Ultraspherical(1,Segment(0.0,1.0))
        @test domainspace(Conversion(S,P)) ==
        ApproxFunBase.SumSpace(ApproxFunBase.HeavisideSpace([-1.0,0.0,1.0]),
            ContinuousSpace(ApproxFunBase.PiecewiseSegment([-1.0,0,1])))

        D = ApproxFunBase.DiracSpace([-1.0,0.0,1.0])
        S2 = ApproxFunBase.SumSpace(D , P)
        f = Fun(S, randn(100))
        @test (Conversion(S,S2) * f)(0.1) ≈ f(0.1)
    end

    @testset "real line integral" begin
        f = x -> x^2/(x^4+1)

        z₁,z₂,z₃,z₄ = exp(im*π/4), exp(3im*π/4), exp(5im*π/4), exp(7im*π/4)

        res₁ = z₁^2 / ((z₁ - z₂)*(z₁ - z₃)*(z₁ - z₄) )
        res₂ = z₂^2 / ((z₂ - z₁)*(z₂ - z₃)*(z₂ - z₄) )

        @test_skip sum(Fun(f, Line())) ≈ 2π*im*(res₁ + res₂)
    end

    @testset "Piecewise Legendre" begin
        f = Fun(x -> exp(-40(x-0.1)^2), Legendre(0..1), 1000)
        @test f(0.2) ≈ exp(-40(0.1)^2)

        f = Fun(x -> exp(-40(x-0.1)^2), Legendre(0..1))
        @test f(0.2) ≈ exp(-40(0.1)^2)

        sp = Legendre(Segment(0 , -1)) ⊕ Legendre(0 .. 1)
        @time f = Fun(x->sign(x)*exp(-40(x-0.1)^2), sp)
        @test f(0.2) ≈ exp(-40(0.1)^2)
        @test f(-0.2) ≈ -exp(-40(0.3)^2)
    end

    @testset "Ambiguous evaluate" begin
        f = Fun(1.0, Chebyshev(NaN..NaN))
        @test f(0.1) ≈ 1.0
    end

    @testset "x + abs(x) (#637)" begin
        x = Fun()
        g = abs(x) + x
        @test ncoefficients(g) == 3
        @test g(0.1) ≈ 0.2
        @test g(-0.2) ≈ 0.0
    end

    @testset "remove point" begin
        f = Fun(identity, (-1..1) \ 0)
        @test f(0.1) ≈ 0.1
        @test f(-0.1) ≈ -0.1
    end
end


@testset "Fourier Shear" begin

    @testset "Complex and real shear produce similar results" begin
        function f(a, b, Δ)
            x = randn((30, 24, 13))
            xc = 0im .+ x
            xc2 = 1im  .* x
            @test shear(x, Δ, a, b) ≈ real(shear(xc, Δ, a, b))
            @test shear(x, Δ, a, b) ≈ imag(shear(xc2, Δ, a, b))
        end

        f(2, 3, 123.1)
        f(3, 2, 13.1)
        f(1, 2, 13.1)
        f(3, 1, 13.1)
    end

    @testset "Test that in-place works in-place" begin
        function f(a, b, Δ)
            x = randn((30, 24, 13))
            xc = randn(ComplexF32, (30, 24, 13))
            xc2 = 1im  .* x
            @test shear!(x, Δ, a, b) ≈ x
            @test shear!(xc, Δ, a, b) ≈ xc
            @test shear!(xc2, Δ, a, b) ≈ xc2
        end

        f(2, 3, 123.1)
        f(3, 2, 13.1)
        f(1, 2, 13.1)
        f(3, 1, 13.1)
    end

    @testset "Fix Nyquist" begin
        @test shear(shear([1 2; 3 4.0], 0.123), -0.123, fix_nyquist = true) == [1.0 2.0; 3.0 4.0]
        @test shear(shear([1 2; 3 4.0], 0.123), -0.123, fix_nyquist = false) != [1.0 2.0; 3.0 4.0]

    end

end


## Test implementing some math logic

# The other trig functions are handled by sympy.
# We find than Cos(Pi) is about 50 time faster than Sin(Pi), etc. in some tests
# So it would be good to implement more of this in Julia

# sympy does more at the moment.
@mkapprule Cos :nargs => 1

do_Cos(mx::Mxpr{:Cos},x) = cos_one_arg(mx,x)

# This is overwritten now anyway
# This does not work. But, If I eval it after starting SJulia, it does work
# Should dispatch on number of args, anyway.
# apprules(mx::Mxpr{:Sin}) = length(mx) == 1 ? sin_one_arg(mx,margs(mx)...) : mx

Cos_pi_coeff(mx::Mxpr{:Cos},c::T) where {T<:Integer} = iseven(c) ? 1 : -1

## Do multiples of pi/2, pi/3, and pi/4
function Cos_pi_coeff(mx::Mxpr{:Cos},c::Rational{T}) where T<:Integer
    n = c.num
    d = c.den
    d == 2 && return zero(d)
    if d == 3
        md = mod(n,6)
        if (md == 2 || md == 4)
            nmx = -1//2
        else
            nmx = 1//2
        end
        return nmx
    end
    if d == 4
        md = mod(n,8)
        if (md == 3 || md == 5)
            nmx = _moosq2
        else
            nmx = _oosq2
        end
        return nmx
    end
    return mx
end

Cos_pi_coeff(mx::Mxpr{:Cos},c::T) where {T<:AbstractFloat} = cospi(c)
Cos_pi_coeff(mx::Mxpr{:Cos},c) = mx

function cos_one_arg(mx::Mxpr{:Cos}, arg::Mxpr{:Times})
    if is_Complex(arg[1])  # find a good way to dispatch on type here.
        (r,i) = reim(arg[1])
        if r == 0  #  Cos(I x) -> Cosh(x)
            nargs = copy(margs(arg))
            nargs[1] = i
            return mxpr(:Cosh,mxpr(:Times,nargs))
        end
    end
    return length(arg) == 2 ? Cos_factor_arg(mx,arg.args...) :
    mx
end

# function cos_one_arg2{T<:Real}(mx::Mxpr{:Cos}, arg::Mxpr{:Times}, arg1::Complex{T})
# end
# function cos_one_arg2(mx::Mxpr{:Cos}, arg::Mxpr{:Times}, arg1)
# end

function Cos_factor_arg(mx::Mxpr{:Cos},f1::T,f2::SJSym) where T<:Number
    if f2 == :Pi
        return Cos_pi_coeff(mx,f1)
    else
        return mx
    end
end

Cos_factor_arg(mx::Mxpr{:Cos},f1,f2) = mx
@inline cos_one_arg(mx::Mxpr{:Cos},arg::Symbol) = arg == :Pi ? -1 : mx
@inline cos_one_arg(mx::Mxpr{:Cos},arg::T) where {T<:Integer} = arg == 0 ? 1 : mx
@inline cos_one_arg(mx::Mxpr{:Cos},x::T) where {T<:AbstractFloat} = cos(x)
@inline cos_one_arg(mx::Mxpr{:Cos},x::Complex{T}) where {T<:AbstractFloat} = cos(x)
@inline cos_one_arg(mx::Mxpr{:Cos},x::Mxpr{:ArcCos}) = length(x) == 1 ? x[1] : mx
function cos_one_arg(mx::Mxpr{:Cos},x::Mxpr{:ArcSin})
    res = mpow((1-mpow(x[1],2)),(1//2))
    return res
end
cos_one_arg(mx::Mxpr{:Cos},x) = mx

#### Sin

# This is overwritten somewhere. Does not work

sin_one_arg(mx::Mxpr{:Sin}, arg::T) where {T<:Integer} = arg == 0 ? 0 : mx

function sin_one_arg(mx::Mxpr{:Sin},arg::Mxpr{:Times})
    if is_Complex(arg[1])  # use dispatch!
        (r,i) = reim(arg[1])
        if r == 0  #  Sin(I x) -> I Sinh(x)
            nargs = copy(margs(arg))
            nargs[1] = i
            return mxpr(:Times, complex(0,1), mxpr(:Sinh,mxpr(:Times,nargs)))
        end
    end
#    todo
    mx
end

sin_one_arg(mx::Mxpr{:Sin},x) = mx

####


using TargetedLearning.Common

facts("Testing automatic influence curves") do
    #simulate scalar estimates for EY1, EY0
    srand(1234)
    n = 50
    a = rand(n)
    b = rand(n)
    ey1 = ScalarEstimate(mean(a), a .- mean(a), "EY1")
    ey0 = ScalarEstimate(mean(b), b .- mean(b), "EY0")
    @fact (ey1 - ey0).ic --> roughly(ey1.ic .- ey0.ic)

    #see http://onlinelibrary.wiley.com/doi/10.1002/sim.3445/abstract appendix
    # for influence curves for logRR and logOR
    @fact log(ey1/ey0).ic --> roughly((1/ey1.psi) .* ey1.ic .- (1/ey0.psi) .* ey0.ic)

    logor = log((ey1/(1-ey1))/(ey0/(1-ey0)))
    @fact logor.ic --> roughly((1/ey1.psi + 1/(1-ey1.psi)) .* ey1.ic - (1/ey0.psi + 1/(1-ey0.psi)) .* ey0.ic)

    logor2 = log(ey1) - log(1-ey1) - log(ey0) + log(1-ey0)
    @fact logor.ic --> roughly(logor2.ic)

    ey1tx1 = ey1 * ey1 * ey1 / sqrt(ey1)
    ey1tx2 = ey1 ^ 2.5
    @fact ey1tx1.psi --> roughly(ey1tx2.psi)
    @fact ey1tx1.ic --> roughly(ey1tx2.ic)
end


# --------10--------20--------30--------40--------50--------60--------70--------80--------90--------100-------110-------120--------130
module AtomicUnits
using PhysicalConstants
import PhysicalConstants.CODATA2018: c_0, ε_0, μ_0, e, a_0, ħ, h, m_e
export fac_i2e, fac_vnm2e, fac_wlenev, fac_au2eV, fac_vm2e
export i2e, e2i, vnm2e, e2vnm, i2vnm, vnm2i
export au2as, as2au, wlen2au, au2wlen, wlen2eV, eV2wlen, eV2au, au2eV
export A2i, i2A
export Up
export au_mass, au_length, au_charge, au_action, au_energy
"""
mass (electron mass)
1 a.u. = mₑ kg = 9.10938e-31 kg
"""
const au_mass = m_e.val
"""
length (Bohr radius)
1 a.u. = a₀ m = 5.291772e-11 m
"""
const au_length = a_0.val
"""
charge (elementary charge)
1 a.u. = e C = 1.602176634e-19 C (exact)
"""
const au_charge = e.val
"""
action (Dirac's constant)
1 a.u. = ħ J⋅s = 1.054571817e-34 J⋅s (exact)
"""
const au_action = ħ.val
"""
energy (Hartree)
1 a.u. = e²/(4πε₀a₀) J
"""
const au_energy = (e^2 / (4π * ε_0 * a_0)).val
"""
fac_vm2e = 5.14221e11
convert factor: electric field [V/nm] ↔ electric field [a.u.]
1 a.u. = Eh/ea₀ V/m = 5.14221e11 V/m
"""
const fac_vm2e = au_energy / (au_charge * au_length)
"""
fac_i2e = 3.50944506e16
convert factor: intensity [W/cm²] ↔ electric field [a.u.]
1 a.u. = ε_0.val * c_0.val * fac_vm2e^2 / 2 / 1e4 = 3.50944506e16 W/cm²
"""
const fac_i2e = ε_0.val * c_0.val * fac_vm2e^2 / 2 / 1e4
"""
fac_vnm2e = 514.221
convert factor: electric field [V/nm] ↔ electric field [a.u.]
1 a.u. = fac_vm2e / 1e9 = 514.221 V/nm
"""
const fac_vnm2e = fac_vm2e / 1e9
"""
fac_au2as = 24.1899
convert factor: time [a.u.] ↔ time [atto. sec]
1 a.u. = ħ/Eh s = 24.1899 as
"""
const fac_au2as = (au_action / au_energy) * 1e18
"""
fac_wleneV = 1239.84190
convert factor: wavelength [nm] ↔ photon energy [eV]
1 nm * 1 eV = hc/e m⋅J = 1239.84190 nm⋅eV
"""
const fac_wleneV = (c_0 * h / e).val * 1e9
"""
fac_au2eV = 27.211386245
convert factor: photon energy [a.u.] ↔ photon energy [eV]
1 a.u. = Eh/e eV = 27.211386245 eV
"""
const fac_au2eV = au_energy / au_charge
## convert functions
# laser amplitude
i2e(fint) = sqrt(fint / fac_i2e)
e2i(famp) = famp^2 * fac_i2e
vnm2e(famp_vnm) = famp_vnm / fac_vnm2e
e2vnm(famp_au) = famp_au * fac_vnm2e
vnm2i(famp_vnm) = e2i(vnm2e(famp_vnm))
i2vnm(fint) = e2vnm(i2e(fint))
"""
Ponderomotive energy (intensity [W/cm²], wavelength [nm]) → Up [a.u.]
"""
Up(fint, wlen_nm) = i2e(fint)^2 / (4 * wlen2au(wlen_nm)^2)
"""
convert (vector potential [a.u.], wavelength [nm]) → intensity [W/cm²]
"""
A2i(A, wlen) = e2i(A * wlen2au(wlen))
"""
convert (intensity [W/cm²], wavelength [nm]) → vector potential [a.u.]
"""
i2A(fint, wlen) = i2e(fint) / wlen2au(wlen)
## time
au2as(t_au) = fac_au2as * t_au
as2au(t_as) = t_as / fac_au2as
## wavelength and photon energy
wlen2au(wlen_nm) = wlen2eV(wlen_nm) / fac_au2eV
au2wlen(freq_au) = fac_au2eV * freq_au
wlen2eV(wlen_nm) = fac_wleneV / wlen_nm
eV2wlen(ene_eV) = fac_wleneV / ene_eV
au2eV(ene_au) = fac_au2eV * ene_au
eV2au(ene_eV) = ene_eV / fac_au2eV
"""
convert wavelength [nm] → period [as]
"""
wlen2period(wlen_nm) = au2as(2π / wlen2au(wlen_nm))
"""
convert energy [eV] → energy [J]
"""
eV2J(ene_eV) = au_charge * ene_eV
"""
convert energy [J] → energy [eV]
"""
J2eV(ene_J) = ene_J / au_charge
end


# LLR, "learning with linear rewards".
# Based on https://arxiv.org/abs/1011.4748.

mutable struct LLR <: Policy end

struct LLRDetails <: PolicyDetails
  solver_time::Float64
end

function choose_action(instance::CombinatorialInstance{T}, policy::LLR, state::State{T}; with_trace::Bool=false) where T
  nArms = length(state.arm_counts)
  if any(v == 0 for v in values(state.arm_counts))
    # Initialisation step: for each random variable (i.e. reward source), try arms, preferring those that have never been tested.
    weights = Dict(arm => (state.arm_counts[arm] == 0.0) ? 1.0 : 0.0 for arm in keys(state.arm_counts))
  else
    # Determine the weights for each arm and use them to solve the combinatorial problem. Parameter required for LLR: maximum solution length.
    # TODO: Find a way to memorise this L! (Called m in the paper, notice...) Maybe within the instance? Or add a method to compute it from the instance (and cache it as soon as it's computed)?
    L_weights = Dict(arm => 1. for arm in keys(state.arm_counts))
    L_sol = solve_linear(instance, L_weights)
    L = length(L_sol)

    weights = Dict(arm => state.arm_average_reward[arm] + sqrt(((L + 1) * log(state.round)) / state.arm_counts[arm]) for arm in keys(state.arm_counts))
  end

  t0 = time_ns()
  sol = solve_linear(instance, weights)
  t1 = time_ns()

  if with_trace
    return sol, LLRDetails((t1 - t0) / 1_000_000_000)
  else
    return sol
  end
end


#  Copyright 2016, Iain Dunning, Joey Huchette, Miles Lubin, and contributors
#  This Source Code Form is subject to the terms of the Mozilla Public
#  License, v. 2.0. If a copy of the MPL was not distributed with this
#  file, You can obtain one at http://mozilla.org/MPL/2.0/.
#############################################################################
# JuMP
# An algebraic modelling langauge for Julia
# See http://github.com/JuliaOpt/JuMP.jl
#############################################################################
# Hock-Schittkowski Nonlinear Test Suite  -  HS109
# This file is JuMP implementation of the model described in
#  W. Hock, K. Schittkowski, Test Examples for Nonlinear Programming
#  Codes, Lecture Notes in Economics and Mathematical Systems,
#  Springer, No, 187, 1981
# More information, including original model description, at
# http://www.ai7.uni-bayreuth.de/downloads.htm
#
# Nonconvex objective with cubics, and constraints contain trignometry and
# products of variables.
#############################################################################

using JuMP
using Base.Test

let

a  = 50.176
b1 = 0.25
b = sin(b1)
c = cos(b1)

L = [0.0, 0.0, -0.55, -0.55, 196, 196, 196, -400, -400]
U = [Inf, Inf,  0.55,  0.55, 252, 252, 252,  800,  800]

m = Model()
@variable(m, L[i] <= x[i=1:9] <= U[i], start = 0.0)

@NLobjective(m, Min, 3 * x[1] + 1e-6 * x[1]^3 + 2 * x[2] + .522074e-6 * x[2]^3)

@NLconstraint(m, x[4] - x[3] + 0.55 >= 0)
@NLconstraint(m, x[3] - x[4] + 0.55 >= 0)
@NLconstraint(m, 2250000 - x[1]^2 - x[8]^2 >= 0)
@NLconstraint(m, 2250000 - x[2]^2 - x[9]^2 >= 0)
@NLconstraint(m,
    x[5] * x[6] * sin(-x[3] - .25) + x[5] * x[7] * sin(-x[4] - .25) +
    2 * b * x[5]^2 - a * x[1] + 400 * a == 0)
@NLconstraint(m,
    x[5] * x[6] * sin(x[3] - .25) + x[6] * x[7] * sin(x[3] - x[4] - .25) +
    2 * b * x[6]^2 - a * x[2] + 400 * a == 0)
@NLconstraint(m,
    x[5] * x[7] * sin(x[4] - .25) + x[6] * x[7] * sin(x[4] - x[3] - .25) +
    2 * b * x[7]^2 + 881.779 * a == 0)
@NLconstraint(m,
    a * x[8] + x[5] * x[6] * cos(-x[3] - .25) +
    x[5] * x[7] * cos(-x[4] - .25) - 200 * a - 2 * c * x[5]^2 +
    .7533e-3 * a * x[5]^2 == 0)
@NLconstraint(m,
    a * x[9] + x[5] * x[6] * cos(x[3] - .25) +
    x[6] * x[7] * cos(x[3] - x[4] - .25) - 2 * c * x[6]^2 +
    .7533e-3 * a * x[6]^2 - 200 * a == 0)
@NLconstraint(m,
    x[5] * x[7] * cos(x[4] - .25) + x[6] * x[7] * cos(x[4] - x[3] - .25) -
    2 * c * x[7]^2 + 22.938 * a + .7533e-3 * a * x[7] ^2 == 0)

solve(m)

@test_approx_eq_eps getobjectivevalue(m) 5326.851310161077 1e-5

end


function backward_rand_kernel(x, Gf, Gpred, Phi) where T
    xf, Pf = meancov(Gf)
    xpred, Ppred = meancov(Gpred)

    J = Pf*Phi'/Ppred
    xs = xf +  J*(x - (Phi*xf  + b))
    Ps = Pf - J*Ppred*J' # as Ppred = M.Phi*P*M.Phi' + M.Q
    rand(Gaussian(xs, Ps))
end

struct BackwardRand{T,S}
    iter::T # iterate backwards through
    M::S
end

function dyniterate(B::BackwardRand, (x, state))
    ϕ = dyniterate(B.iter, state)
    ϕ === nothing && return nothing
    (n, U), state = ϕ
    Gf, Gpred = U[1], U[2]
    Gs = backward_rand_kernel(x, Gf, Gpred, Phi(RTS.M.sys, (n, U)))
    n=>x, (x, state)
end


using LinearAlgebra
using BlockFactorizations
using BenchmarkTools

d = 512
n, m = 16, 8
A = [randn(d, d) for i in 1:n, j in 1:m]
B = BlockFactorization(A)

x = [randn(d) for _ in 1:m]
y = [zeros(d) for _ in 1:n]
@btime A*x
@btime mul!(y, A, x);

x = randn(d*m)
y = zeros(d*n)
@btime B*x;
@btime mul!(y, B, x);

d = 512
n, m = 16, 8
# testing with diagonal components
A = [Diagonal(randn(d)) for i in 1:n, j in 1:m]
B = BlockFactorization(A)

x = [randn(d) for _ in 1:m]
y = [zeros(d) for _ in 1:n]
@btime A*x
@btime mul!(y, A, x);

x = randn(d*m)
y = zeros(d*n)
@btime B*x;
@btime mul!(y, B, x);


module LeastAction

using FindSteadyStates
using SteadyStateDiffEq
using LinearAlgebra
using ModelingToolkit
using ForwardDiff
using Optim
using NLsolve
using Parameters

include("utils.jl")
include("fixedpoint_gen.jl")
include("DiffusionMatrix.jl")
include("gradient_gen.jl")
include("action.jl")

export fixedpoint_gen
export gradient_gen
export action

end


"""
estimate(x; β::Real, δ::Real)

Estimate a stochastic volatility model. Volatility is modelled as a random walk.

The hyperparameters 2/3 < β < 1 and 0 < δ ≤ 1 are discount factors, controlling the shocks to the Σ and Ω respectively.
"""

function estimate(ssm::StateSpaceModel)
    y, F, G, m, P, S, β, δ, Δ, ν, n, k = ssm.y, ssm.F, ssm.G, ssm.m, ssm.P, ssm.S, ssm.β, ssm.δ, ssm.Δ, ssm.ν, ssm.k, ssm.k

    T, p = size(y)
    d = size(m, 1)

    predicted = (μ = zeros(T, p), Σ = zeros(T, p, p), e = zeros(T, p), u = zeros(T, p), m = zeros(T, d, p), R = zeros(T, d, d), S = zeros(T, p, p))
    filtered  = (μ = zeros(T, p), Σ = zeros(T, p, p), e = zeros(T, p), u = zeros(T, p), m = zeros(T, d, p), P = zeros(T, d, d), S = zeros(T, p, p))

    for t = 1:T
        # Predict at time t-1
        R = get_R(P, G, Δ)
        Q = get_Q(F, R)
        μ = m' * G' * F
        e = y[t, :] - μ
        Σ = Q * (1 - β) / (3β*k - 2k) * S
        u = inv(cholesky(Σ).L) * e

        predicted.μ[t, :]    = μ
        predicted.Σ[t, :, :] = Σ
        predicted.e[t, :, :] = e
        predicted.u[t, :, :] = u
        predicted.m[t, :, :] = m
        predicted.R[t, :, :] = R
        predicted.S[t, :, :] = S

        # Update at time t
        K = get_K(F, R, Q)
        m = get_m(m, G, K, e)
        P = get_P(R, K, Q)
        S = get_S(S, k, e, Q)
        μ = m' * F
        e = y[t, :] - μ
        Σ = Q * (1 - β) / (2β - 1) * S
        u = inv(cholesky(Σ).L) * e

        filtered.μ[t, :]    = μ
        filtered.Σ[t, :, :] = Σ
        filtered.e[t, :, :] = e
        filtered.u[t, :, :] = u
        filtered.m[t, :, :] = m
        filtered.P[t, :, :] = P
        filtered.S[t, :, :] = S
    end
    return (predicted = predicted, filtered = filtered)
end


using Revise
using AdFem
using PyCall
using LinearAlgebra
using PyPlot
using SparseArrays
using MAT
np = pyimport("numpy")

stepsize = 1
if length(ARGS)==1
  global stepsize = parse(Int64, ARGS[1])
end
@info stepsize

## alpha-scheme
β = 1/4; γ = 1/2
a = b = 0.1    # damping

mmesh = Mesh(10, 10, 1/10, degree=1)
NT = 500
Δt = 2.0/NT
ηmax = 1
ηmin = 0.5

bdedge = bcedge((x1, y1, x2, y2)->(x1>1-1e-5 && x2>1-1e-5), mmesh)
bdnode = bcnode((x1, y1)->(y1<1e-5), mmesh)

invη = 50 * constant(ones(get_ngauss(mmesh)))
μ = 1.0
λ = 1.0

fn_G = invη->begin
  G = tensor([1/Δt+2/3*μ*invη -μ/3*invη 0.0
    -μ/3*invη 1/Δt+2/3*μ*invη 0.0
    0.0 0.0 1/Δt+μ*invη])
  invG = inv(G)
end
invG = map(fn_G, invη)
S = tensor([2μ/Δt+λ/Δt λ/Δt 0.0
    λ/Δt 2μ/Δt+λ/Δt 0.0
    0.0 0.0 μ/Δt])
H = invG*S

M = compute_fem_mass_matrix1(mmesh)
Zero = spzero(mmesh.ndof, mmesh.ndof)
M = [M Zero;Zero M]

K = compute_fem_stiffness_matrix(H, mmesh)
C = a*M + b*K # damping matrix
L = M + γ*Δt*C + β*Δt^2*K
bddof = [bdnode; bdnode .+ mmesh.ndof]
L, _ = impose_Dirichlet_boundary_conditions(L, zeros(2mmesh.ndof), bddof, zeros(length(bddof)))
L = factorize(L)

a = TensorArray(NT+1); a = write(a, 1, zeros(2mmesh.ndof))
v = TensorArray(NT+1); v = write(v, 1, zeros(2mmesh.ndof))
d = TensorArray(NT+1); d = write(d, 1, zeros(2mmesh.ndof))
U = TensorArray(NT+1); U = write(U, 1, zeros(2mmesh.ndof))
Sigma = TensorArray(NT+1); Sigma = write(Sigma, 1, zeros(get_ngauss(mmesh), 3))
Varepsilon = TensorArray(NT+1); Varepsilon = write(Varepsilon, 1, zeros(get_ngauss(mmesh), 3))

Forces = zeros(NT, 2mmesh.ndof)
T = eval_f_on_boundary_edge((x,y)->0.1, bdedge, mmesh)
T = [-T T]
rhs = compute_fem_traction_term(T, bdedge, mmesh)
for i = 1:NT
  Forces[i, :] = rhs
end
Forces = constant(Forces)

function condition(i, tas...)
  i <= NT
end

function body(i, tas...)
  a_, v_, d_, U_, Sigma_, Varepsilon_ = tas
  a = read(a_, i)
  v = read(v_, i)
  d = read(d_, i)
  U = read(U_, i)
  Sigma = read(Sigma_, i)
  Varepsilon = read(Varepsilon_, i)

  res = batch_matmul(invG/Δt, Sigma)
  F = compute_strain_energy_term(res, mmesh) - K * U
  rhs = Forces[i] - F

  td = d + Δt*v + Δt^2/2*(1-2β)*a
  tv = v + (1-γ)*Δt*a
  rhs = rhs - C*tv - K*td
  rhs = scatter_update(rhs, bddof, zeros(length(bddof)))

  ## alpha-scheme
  a = L\rhs # bottleneck
  d = td + β*Δt^2*a
  v = tv + γ*Δt*a
  U_new = d

  Varepsilon_new = eval_strain_on_gauss_pts(U_new, mmesh)
  Sigma_new = update_stress_viscosity(Varepsilon_new, Varepsilon, Sigma, invη, μ*ones(get_ngauss(mmesh)), λ*ones(get_ngauss(mmesh)), Δt)

  i+1, write(a_, i+1, a), write(v_, i+1, v), write(d_, i+1, d), write(U_, i+1, U_new),
        write(Sigma_, i+1, Sigma_new), write(Varepsilon_, i+1, Varepsilon_new)
end

i = constant(1, dtype=Int32)
_, _, _, _, u, sigma, varepsilon = while_loop(condition, body,
                  [i, a, v, d, U, Sigma, Varepsilon])

U = stack(u)
Sigma = stack(sigma)
Varepsilon = stack(varepsilon)

sess = Session(); init(sess)
Uval,Sigmaval, Varepsilonval = run(sess, [U, Sigma, Varepsilon])

# p = visualize_von_mises_stress(Sigmaval[1:5:end,:,:], m, n, h); saveanim(p, "space_s.gif")
p = visualize_displacement(Uval[1:5:end,:], mmesh)
saveanim(p, "space_u.gif")

# visualize_von_mises_stress(Sigmaval[end,:,:], m, n, h); savefig("space_s.pdf")
# visualize_displacement(Uval[end,:], m, n, h); savefig("space_u.pdf")

# visualize_inv_eta(run(sess, invη), "true")


module QuantumRelay

using SymPy
using Iterators
using PolyLib
using JuMP
using Clp
using LLLplus
using GSL
using PyPlot
using Lora
using Distributions

export
	qrelay_op,
	op_mat,
	scan_maker,
	QRelaySampler

include("util.jl")
include("operators.jl")
include("scan.jl")
include("distributions.jl")

type QRelaySampler
	prob::Function
	psetproposal::Function

	function QRelaySampler{T<:Int}(mat::Array{T, 2}, coef, omega, pdet0)
		U, S, V = PolyLib.smith_normal_form(mat)
		Ui = PolyLib.inverse(U)
		Vi = PolyLib.inverse(V)
		s = diag(S)
		r = countnz(s)
		s0 = s[1:r]
		@assert s0 == ones(r)
		ui1 = Ui[1:r, :]
		ui2 = Ui[r+1:end, :]
		vi1 = Vi[:, 1:r]
		vi2 = Vi[:, r+1:end]
		vi2 = lll(vi2)[1]
		T0 = vi1*ui1
		ui2oc = orthocomp(ui2)
		setc, scan = scan_maker(vi2)

		function prob(na)
		    @assert countnz(ui2*na) == 0
		    b = T0*na
		    setc(-b)
		    total = 0.0
		    for x in Task(scan)
		        nab = vi2*x + b
		        total += prod([c.^complex(n)/factorial(n) for (c, n) in zip(coef, nab)])
		    end
		    return abs(total*omega)^2
		end

		function prob(q, na, mask)
		    q0 = round(Int, q.>0)
		    m0 = round(Int, mask)
		    return prod((q0 + (1-2q0).*pdet0(na)).^m0)
		end

		psetproposal(x::Vector) = QuantumRelay.OrthoNNDist(x, ui2oc)

		new(prob, psetproposal)
	end
end

end


"""
    lefttransfer(HL₀::AbstractMatrix, hL::Union{Nothing, AbstractPiecewise},
                    Ψ₁::FiniteCMPS, Ψ₂::FiniteCMPS = Ψ₁;
                    Kmax = 50, tol = defaulttol(Ψ₁))

Solve the non-homogeneous (or homogeneous if hL === nothing) Lindblad equation

```math
∂ₓ(HL) = Q₂' * HL + HL * Q₁ + R₂' * HL * R₁ + hL
```

with `Q₁` and `R₁` the matrices of the finite cMPS `Ψ₁` (similar for `Q₂` and `R₂`, where
`Ψ₂ = Ψ₁` by default), over the whole domain `(a,b) = domain(Ψ₁)`, starting from the initial
condition `HL(a) = HL₀`, all the way up to `b`. The states `Ψ₁` and `Ψ₂` are assumed to have
equal domain and grid; some checks are in place.

The result is given by

```math
⟨HL(x)| = ⟨HL₀|exp(∫ₐˣ T(z) dz) + ∫ₐˣ [ ⟨hL(u)|exp(∫ᵤˣ T(z) dz) ] du
```

and returned as `HL, converged` with `HL::Piecewise` and where `converged::Bool` is `true`
or `false` depending on convergence up to the requested tolerance.
"""
function lefttransfer(HL₀::AbstractMatrix, hL::Union{Nothing, AbstractPiecewise},
                        Ψ₁::FiniteCMPS, Ψ₂::FiniteCMPS = Ψ₁;
                        Kmax = 50, tol = defaulttol(Ψ₁))

    Q₁ = Ψ₁.Q
    R₁s = Ψ₁.Rs
    Q₂ = Ψ₂.Q
    R₂s = Ψ₂.Rs
    grid = nodes(Q₁)
    N = length(grid) - 1
    a = first(grid)
    b = last(grid)
    if Ψ₁ !== Ψ₂
        grid == nodes(Q₂) || throw(DomainMismatch())
    end
    if !isnothing(hL)
        grid == nodes(hL) || throw(DomainMismatch())
    end
    HLᵢ = TaylorSeries([HL₀], a)
    HLs = Vector{typeof(HLᵢ)}(undef, N)
    HLs[1] = HLᵢ
    infoL = true
    for i = 1:N
        HLᵢ = HLs[i]
        xᵢ = grid[i]
        xⱼ = grid[i+1]
        Δxᵢ = grid[i+1] - grid[i]
        Q₁ᵢ = shift!(Q₁[i], xᵢ)
        R₁ᵢs = shift!.(getindex.(R₁s, i), xᵢ)
        Q₂ᵢ = shift!(Q₂[i], xᵢ)
        R₂ᵢs = shift!.(getindex.(R₂s, i), xᵢ)
        hLᵢ = isnothing(hL) ? nothing : shift(hL[i], xᵢ)

        # build Taylor coefficients (i.e. solve triangular problem)
        HLᵢ, converged = _lefttransfer!(HLᵢ, hLᵢ, Q₁ᵢ, R₁ᵢs, Δxᵢ, Q₂ᵢ, R₂ᵢs;
                                        Kmax = Kmax, tol = tol)
        infoL &= converged
        shift!(HLᵢ, (xᵢ+xⱼ)/2)

        # initialize next ρ element
        if i < N
            HLs[i+1] = TaylorSeries([HLᵢ(xⱼ)], xⱼ)
        end
    end
    HL = Piecewise(grid, HLs)
    return HL, infoL
end

"""
    righttransfer(HR₀::AbstractMatrix, hR::Union{Nothing, AbstractPiecewise},
                    Ψ₁::FiniteCMPS, Ψ₂::FiniteCMPS = Ψ₁;
                    Kmax = 50, tol = defaulttol(Ψ₁))

Solve the non-homogeneous (or homogeneous if hR === nothing) Lindblad equation

```math
-∂ₓ(HR) = Q₁ * HR + HR * Q₂' + R₁ * HR * R₂ + hR
```

with `Q₁` and `R₁` the matrices of the finite cMPS `Ψ₁` (similar for `Q₂` and `R₂`, where
`Ψ₂ = Ψ₁` by default), over the whole domain `(a,b) = domain(Ψ₁)`, starting from the initial
condition `HR(b) = HR₀`, all the way down to `a`. The states `Ψ₁` and `Ψ₂` are assumed to
have equal domain and grid; some checks are in place.

The result is given by

```math
|HR(x)⟩ = exp(∫ₓᵇ T(z) dz)|HR₀⟩ + ∫ₓᵇ exp(∫ₓᵘ T(z) dz)|hR(u)⟩ du
```

and returned as `HL, converged` with `HL::Piecewise` and where `converged::Bool` is `true`
or `false` depending on convergence up to the requested tolerance.
"""
function righttransfer(HR₀::AbstractMatrix, hR::Union{Nothing, AbstractPiecewise},
                        Ψ₁::FiniteCMPS, Ψ₂::FiniteCMPS = Ψ₁;
                        Kmax = 50, tol = defaulttol(Ψ₁))

    Q₁ = Ψ₁.Q
    R₁s = Ψ₁.Rs
    Q₂ = Ψ₂.Q
    R₂s = Ψ₂.Rs
    grid = nodes(Q₁)
    N = length(grid) - 1
    a = first(grid)
    b = last(grid)
    if Ψ₁ !== Ψ₂
        grid == nodes(Q₂) || throw(DomainMismatch())
    end
    if !isnothing(hR)
        grid == nodes(hR) || throw(DomainMismatch())
    end
    HRᵢ = TaylorSeries([HR₀], b)
    HRs = Vector{typeof(HRᵢ)}(undef, N)
    HRs[N] = HRᵢ
    infoR = true
    for i = N:-1:1
        HRᵢ = HRs[i]
        xᵢ = grid[i]
        xⱼ = grid[i+1]
        Δxᵢ = grid[i+1] - grid[i]
        Q₁ᵢ = shift!(Q₁[i], xⱼ)
        R₁ᵢs = shift!.(getindex.(R₁s, i), xⱼ)
        Q₂ᵢ = shift!(Q₂[i], xⱼ)
        R₂ᵢs = shift!.(getindex.(R₂s, i), xⱼ)
        hRᵢ = isnothing(hR) ? nothing : shift(hR[i], xⱼ)

        # build Taylor coefficients (i.e. solve triangular problem)
        HRᵢ, converged = _righttransfer!(HRᵢ, hRᵢ, Q₁ᵢ, R₁ᵢs, Δxᵢ, Q₂ᵢ, R₂ᵢs;
                                            Kmax = Kmax, tol = tol)
        infoR &= converged
        shift!(HRᵢ, (xᵢ+xⱼ)/2)

        # initialize next ρ element
        if i > 1
            HRs[i-1] = TaylorSeries([HRᵢ(xᵢ)], xᵢ)
        end
    end
    HR = Piecewise(grid, HRs)
    return HR, infoR
end

"""
    _lefttransfer!(ρ::TaylorSeries, h::TaylorSeries,
                    Q₁::TaylorSeries, R₁s::Tuple{Vararg{TaylorSeries}}, Δx,
                    Q₂::TaylorSeries = Q₁, R₂s::Tuple{Vararg{TaylorSeries}} = R₁s;
                    Kmax = 50, tol = defaulttol(ρ))

Solve the non-homogeneous (or homogeneous if h === nothing) Lindblad equation

```math
∂ₓρ = Q₂'*ρ + ρ*Q₁ + R₂'*ρ*R₁ + h
```

in place (i.e. the Taylor coefficients of `ρ` will be inserted), starting from the initial
conditions `ρ[0]` at position `x = offset(ρ)`. The solution will have `degree(ρ) <= Kmax`
and is found up to a tolerance `tol` in the interval `Δx`, i.e. the last few Taylor
coefficients of `ρ` and `h` should satisfy `ρ[k]*Δx^k < tol`. The result is returned as `ρ,
converged`, where `converged::Bool` is `true` or `false` depending on convergence up to the
given tolerance.

Note that the coefficients of `h` are destroyed in the process, if you want to preserve
those you should take a copy beforehand.
"""
function _lefttransfer!(ρ::TaylorSeries, h::Union{Nothing, TaylorSeries},
                            Q₁::TaylorSeries, R₁s::Tuple{Vararg{TaylorSeries}}, Δx,
                            Q₂::TaylorSeries = Q₁, R₂s::Tuple{Vararg{TaylorSeries}} = R₁s;
                            Kmax = 50, tol = defaulttol(ρ))

    Kmin = max(degree(Q₁), degree(Q₂), maximum(degree, R₁s), maximum(degree, R₂s)) + 1
    temp = zero(ρ[0])
    T = eltype(ρ[0])
    for k = 1:Kmax
        ρᵏ = isnothing(h) ? zero(ρ[0]) : h[k-1]
        # solve triangular system:
        for l = 0:min(k-1, degree(Q₁))
            mul!(ρᵏ, ρ[k-1-l], Q₁[l], one(T), one(T))
        end
        for l = 0:min(k-1, degree(Q₂))
            mul!(ρᵏ, Q₂[l]', ρ[k-1-l], one(T), one(T))
        end
        for (R₁, R₂) in zip(R₁s, R₂s)
            for l = 0:min(k-1, degree(R₂))
                for m = 0:min(k-1-l, degree(R₁))
                    mul!(temp, ρ[k-1-l-m], R₁[m], one(T), zero(T))
                    mul!(ρᵏ, R₂[l]', temp, one(T), one(T))
                end
            end
        end
        rmul!(ρᵏ, 1/k)
        ρ[k] = ρᵏ
        # check for convergence
        if k > Kmin
            converged = true
            for l = k-Kmin+1:k
                if norm(ρ[l]) * (Δx^l) > tol
                    converged = false
                end
            end
            if converged
                return ρ, true
            end
        end
    end
    return ρ, false
end

"""
    _righttransfer!(ρ::TaylorSeries, h::TaylorSeries,
                    Q₁::TaylorSeries, R₁s::Tuple{Vararg{TaylorSeries}}, Δx,
                    Q₂::TaylorSeries = Q₁, R₂s::Tuple{Vararg{TaylorSeries}} = R₁s;
                    Kmax = 50, tol = defaulttol(ρ))

Solve the non-homogeneous (or homogeneous if h === nothing) Lindblad equation

```math
∂ₓρ = -(Q₂'*ρ + ρ*Q₁ + R₂'*ρ*R₁) - h
```

in place (i.e. the Taylor coefficients of `ρ` will be inserted), starting from the initial
conditions `ρ[0]` at position `x = offset(ρ)`. The solution will have `degree(ρ) <= Kmax`
and is found up to a tolerance `tol` in the interval `Δx`, i.e. the last few Taylor
coefficients of `ρ` and `h` should satisfy `ρ[k]*Δx^k < tol`. The result is returned as `ρ,
converged`, where `converged::Bool` is `true` or `false` depending on convergence up to the
given tolerance.

Note that the coefficients of `h` are destroyed in the process, if you want to preserve
those you should take a copy beforehand. Also note that all TaylorSeries are assumed to have
the same offset.
"""
function _righttransfer!(ρ::TaylorSeries, h::Union{Nothing, TaylorSeries},
                            Q₁::TaylorSeries, R₁s::Tuple{Vararg{TaylorSeries}}, Δx,
                            Q₂::TaylorSeries = Q₁, R₂s::Tuple{Vararg{TaylorSeries}} = R₁s;
                            Kmax = 50, tol = defaulttol(ρ))

    Kmin = max(degree(Q₁), degree(Q₂), maximum(degree, R₁s), maximum(degree, R₂s)) + 1
    temp = zero(ρ[0])
    T = eltype(ρ[0])
    for k = 1:Kmax
        ρᵏ = isnothing(h) ? zero(ρ[0]) : rmul!(h[k-1], -1)
        # solve triangular system:
        for l = 0:min(k-1, degree(Q₁))
            mul!(ρᵏ, Q₁[l], ρ[k-1-l], -one(T), one(T))
        end
        for l = 0:min(k-1, degree(Q₂))
            mul!(ρᵏ, ρ[k-1-l], Q₂[l]', -one(T), one(T))
        end
        for (R₁, R₂) in zip(R₁s, R₂s)
            for l = 0:min(k-1, degree(R₁))
                for m = 0:min(k-1-l, degree(R₂))
                    mul!(temp, ρ[k-1-l-m], R₂[m]', one(T), zero(T))
                    mul!(ρᵏ, R₁[l], temp, -one(T), one(T))
                end
            end
        end
        rmul!(ρᵏ, 1/k)
        ρ[k] = ρᵏ
        # check for convergence
        if k > Kmin
            converged = true
            for l = k-Kmin+1:k
                if norm(ρ[l]) * (Δx^l) > tol
                    converged = false
                end
            end
            if converged
                return ρ, true
            end
        end
    end
    # not converged:
    return ρ, false
end


module PSDMatrices

import Base: \, /, size, inv, copy, copy!, ==, show, similar, Matrix
using LinearAlgebra
import LinearAlgebra: det, logabsdet, diag

struct PSDMatrix{T,FactorType} <: AbstractMatrix{T}
    R::FactorType
    PSDMatrix(R::AbstractMatrix{T}) where {T} = new{T,typeof(R)}(R)
end

# Base overloads
Matrix(M::PSDMatrix) = M.R' * M.R
size(M::PSDMatrix) = (size(M.R, 2), size(M.R, 2))
inv(M::PSDMatrix) = PSDMatrix(inv(M.R'))
\(A::PSDMatrix, B::AbstractVecOrMat) = A.R \ (A.R' \ B)
/(B::AbstractVecOrMat, A::PSDMatrix) = B / A.R / A.R'
copy(M::PSDMatrix) = PSDMatrix(copy(M.R))
similar(M::PSDMatrix, element_type::Type=eltype(M)) = PSDMatrix(similar(M.R, element_type))
copy!(dst::PSDMatrix, src::PSDMatrix) = (copy!(dst.R, src.R); dst)
==(M1::PSDMatrix, M2::PSDMatrix) = M1.R == M2.R  # todo: same as isequal()?!
function show(io::IO, M::PSDMatrix)
    print(io, "$(size(M,1))x$(size(M,2)) $(typeof(M)); R=")
    show(io, M.R)
end
function show(io::IO, m::MIME"text/plain", M::PSDMatrix)
    println(io, "$(size(M,1))x$(size(M,2)) $(typeof(M)) ")
    print(io, " Right square root: R=")
    show(io, m, M.R)
end

# LinearAlgebra overloads

function det(M::PSDMatrix)
    confirm_factor_is_square(M)
    return det(M.R)^2
end

function logabsdet(M::PSDMatrix)
    confirm_factor_is_square(M)
    _logabsdet, _sign = logabsdet(M.R)
    return 2 * _logabsdet, _sign^2
end

function diag(M::PSDMatrix)
    out = similar(M.R, size(M.R, 2))
    sum!(abs2, out', M.R)
    return out
end

function confirm_factor_is_square(M::PSDMatrix)
    if size(M.R, 1) != size(M.R, 2)
        msg = (
            "The requested operation is not available for a PSDMatrix with a non-square factor." *
            "The factor of the received PSDMatrix has dimensions ($(size(M.R,1)), $(size(M.R,2))). " *
            "Try turning the PSDMatrix into a dense matrix first."
        )
        throw(MethodError(msg))
    end
end

# Custom functions

X_A_Xt(; A::PSDMatrix, X::AbstractMatrix) = PSDMatrix(A.R * X')
X_A_Xt(A::PSDMatrix, X::AbstractMatrix) = X_A_Xt(A=A, X=X)
function X_A_Xt!(out::PSDMatrix; A::PSDMatrix, X::AbstractMatrix)
    mul!(out.R, A.R, X')
    return out
end
X_A_Xt!(out::PSDMatrix, A::PSDMatrix, X::AbstractMatrix) = X_A_Xt!(out, A=A, X=X)

function add_cholesky(A::PSDMatrix, B::PSDMatrix)
    sum_dense = Matrix(A) + Matrix(B)
    factor = cholesky(sum_dense).U
    return PSDMatrix(factor)
end

function add_qr(A::PSDMatrix, B::PSDMatrix)
    stack = vcat(A.R, B.R)
    matrix = PSDMatrix(stack)
    return triangularize_factor(matrix)
end

function triangularize_factor(M::PSDMatrix)
    R = qr(M.R).R
    return PSDMatrix(UpperTriangular(R))
end

export PSDMatrix
export add_cholesky
export add_qr
export triangularize_factor
export X_A_Xt
export X_A_Xt!

end


"""
    resonance_eigenproblem(simulation, k, ka=0, kb=0) -> A,B,σ

generate matrices from `simulation` to solve {∇²+εk²}ψ=0 in the form `A*v=λB*b` with shift `σ`
"""
function resonance_eigenproblem(sim::Simulation, k::Number, ka::Number=0, kb::Number=0)
    L = Laplacian(sim,k,ka,kb)
    return L.D, -L.M, k^2
end

"""
    resonance_eigenproblem(sim, k, ka=0, kb=0; η=0) -> A,B,σ

generate matrices from `simulation` to solve {∇²+εk²}u=-ηFk²u in the form `A*v=λB*b` with shift `σ
"""
function cf_eigenproblem(sim::Simulation, k::Number, ka::Number=0, kb::Number=0; η::Number=0)
    L = Laplacian(sim,k,ka,kb)
    return L.D+L.M*k^2, -L.Mf*k^2, η
end


using InterestRates

# First, create a curve instance.

vert_x = [11, 15, 50, 80] # for interpolation methods, represents the days to maturity
vert_y = [0.10, 0.15, 0.14, 0.17] # yield values

dt_curve = Date(2015,08,03)

mycurve = InterestRates.IRCurve("dummy-simple-linear", InterestRates.Actual365(),
    InterestRates.SimpleCompounding(), InterestRates.Linear(), dt_curve,
    vert_x, vert_y)

# yield for a given maturity date
y = zero_rate(mycurve, Date(2015,08,25))
# 0.148

# forward rate between two future dates
fy = forward_rate(mycurve, Date(2015,08,25), Date(2015, 10, 10))
# 0.16134333771591897

# Discount factor for a given maturity date
df = discountfactor(mycurve, Date(2015,10,10))
# 0.9714060637029466

# Effective Rate Factor for a given maturity
erf = ERF(mycurve, Date(2015,10,10))
# 1.0294356164383562

# Effective Rate for a given maturity
er = ER(mycurve, Date(2015,10,10))
# 0.029435616438356238


# dcontract, dot, tdot, otimes, cross
"""
```julia
dcontract(::SecondOrderTensor, ::SecondOrderTensor)
dcontract(::SecondOrderTensor, ::FourthOrderTensor)
dcontract(::FourthOrderTensor, ::SecondOrderTensor)
dcontract(::FourthOrderTensor, ::FourthOrderTensor)
```
Computes the double contraction between two tensors.
The symbol `⊡`, written `\\boxdot`, is overloaded for double contraction.
The reason `:` is not used is because it does not have the same precedence as multiplication.

**Example:**

```jldoctest
julia> A = rand(SymmetricTensor{2, 2});

julia> B = rand(SymmetricTensor{2, 2});

julia> dcontract(A,B)
1.9732018397544984

julia> A ⊡ B
1.9732018397544984
```
"""
@inline dcontract{dim}(S1::Tensor{2, dim}, S2::Tensor{2, dim}) = tovector(S1) ⋅ tovector(S2)

@inline function dcontract{dim}(S1::Tensor{4, dim}, S2::Tensor{4, dim})
    Tensor{4, dim}(tomatrix(S1) * tomatrix(S2))
end

@inline function dcontract{dim}(S1::Tensor{4, dim}, S2::Tensor{2, dim})
    Tensor{2, dim}(tomatrix(S1) * tovector(S2))
end

@inline function dcontract{dim}(S1::Tensor{2, dim}, S2::Tensor{4, dim})
    Tensor{2, dim}(tomatrix(S2)' * tovector(S1))
end

const ⊡ = dcontract

# Specialized methods for symmetric tensors
@generated function dcontract{dim}(S1::SymmetricTensor{2, dim}, S2::SymmetricTensor{2, dim})
    idx2(i,j) = compute_index(SymmetricTensor{2, dim}, i, j)
    ex1, ex2 = Expr[], Expr[]
    for j in 1:dim, i in j:dim
        if i == j
            push!(ex1, :(get_data(S1)[$(idx2(i, j))]))
            push!(ex2, :(get_data(S2)[$(idx2(i, j))]))
        else
            push!(ex1, :(2 * get_data(S1)[$(idx2(i, j))]))
            push!(ex2, :(    get_data(S2)[$(idx2(i, j))]))
        end
    end
    exp = make_muladd_exp(ex1, ex2)
    return quote
        $(Expr(:meta, :inline))
        $exp
    end
end

@generated function dcontract{dim}(S1::Tensor{2, dim}, S2::SymmetricTensor{4, dim})
    idx4(i,j,k,l) = compute_index(SymmetricTensor{4, dim}, i, j, k, l)
    idx2(i,j) = compute_index(Tensor{2, dim}, i, j)
    exps = Expr(:tuple)
    for l in 1:dim, k in l:dim
        ex1, ex2 = Expr[], Expr[]
        for j in 1:dim, i in 1:dim
            push!(ex1, :(data2[$(idx2(i, j))]))
            push!(ex2, :(data4[$(idx4(i, j, k, l))]))
        end
        push!(exps.args, make_muladd_exp(ex1, ex2))
    end
    quote
        $(Expr(:meta, :inline))
        data2 = get_data(S1)
        data4 = get_data(S2)
        @inbounds r = $exps
        SymmetricTensor{2, dim}(r)
    end
end

@generated function dcontract{dim}(S1::SymmetricTensor{4, dim}, S2::Tensor{2, dim})
    idx4(i,j,k,l) = compute_index(SymmetricTensor{4, dim}, i, j, k, l)
    idx2(k,l) = compute_index(Tensor{2, dim}, k, l)
    exps = Expr(:tuple)
    for j in 1:dim, i in j:dim
        ex1, ex2 = Expr[], Expr[]
        for l in 1:dim, k in 1:dim
            push!(ex1, :(data4[$(idx4(i, j, k, l))]))
            push!(ex2, :(data2[$(idx2(k, l))]))
        end
        push!(exps.args, make_muladd_exp(ex1, ex2))
    end
    quote
        $(Expr(:meta, :inline))
        data2 = get_data(S2)
        data4 = get_data(S1)
        @inbounds r = $exps
        SymmetricTensor{2, dim}(r)
    end
end

@generated function dcontract{dim}(S1::SymmetricTensor{2, dim}, S2::SymmetricTensor{4, dim})
    idx4(i,j,k,l) = compute_index(SymmetricTensor{4, dim}, i, j, k, l)
    idx2(i,j) = compute_index(SymmetricTensor{2, dim}, i, j)
    exps = Expr(:tuple)
    for l in 1:dim, k in l:dim
        ex1, ex2 = Expr[], Expr[]
        for j in 1:dim, i in j:dim
            if i == j
                push!(ex1, :(data2[$(idx2(i, j))]))
                push!(ex2, :(data4[$(idx4(i, j, k, l))]))
            else
                push!(ex1, :(2 * data2[$(idx2(i, j))]))
                push!(ex2, :(    data4[$(idx4(i, j, k, l))]))
            end
        end
        push!(exps.args, make_muladd_exp(ex1, ex2))
    end
    quote
        $(Expr(:meta, :inline))
        data2 = get_data(S1)
        data4 = get_data(S2)
        @inbounds r = $exps
        SymmetricTensor{2, dim}(r)
    end
end

@generated function dcontract{dim}(S1::SymmetricTensor{4, dim}, S2::SymmetricTensor{2, dim})
    idx4(i,j,k,l) = compute_index(SymmetricTensor{4, dim}, i, j, k, l)
    idx2(k,l) = compute_index(SymmetricTensor{2, dim}, k, l)
    exps = Expr(:tuple)
    for j in 1:dim, i in j:dim
        ex1, ex2 = Expr[], Expr[]
        for l in 1:dim, k in l:dim
            if k == l
                push!(ex1, :(data4[$(idx4(i, j, k, l))]))
                push!(ex2, :(data2[$(idx2(k, l))]))
            else
                push!(ex1, :(2 * data4[$(idx4(i, j, k, l))]))
                push!(ex2, :(    data2[$(idx2(k, l))]))
            end
        end
        push!(exps.args, make_muladd_exp(ex1, ex2))
    end
    quote
        $(Expr(:meta, :inline))
        data2 = get_data(S2)
        data4 = get_data(S1)
        @inbounds r = $exps
        SymmetricTensor{2, dim}(r)
    end
end

@generated function dcontract{dim}(S1::SymmetricTensor{4, dim}, S2::SymmetricTensor{4, dim})
    idx4(i,j,k,l) = compute_index(SymmetricTensor{4, dim}, i, j, k, l)
    exps = Expr(:tuple)
    for l in 1:dim, k in l:dim, j in 1:dim, i in j:dim
        ex1, ex2 = Expr[], Expr[]
        for n in 1:dim, m in n:dim
            if m == n
                push!(ex1, :(data1[$(idx4(i, j, m, n))]))
                push!(ex2, :(data2[$(idx4(m, n, k, l))]))
            else
                push!(ex1, :(2 * data1[$(idx4(i, j, m, n))]))
                push!(ex2, :(    data2[$(idx4(m, n, k, l))]))
            end
        end
        push!(exps.args, make_muladd_exp(ex1, ex2))
    end
    quote
        $(Expr(:meta, :inline))
        data2 = get_data(S2)
        data1 = get_data(S1)
        @inbounds r = $exps
        SymmetricTensor{4, dim}(r)
    end
end

# Promotion
@inline dcontract{dim}(S1::Tensor{2, dim}, S2::SymmetricTensor{2, dim}) = dcontract(promote(S1, S2)...)
@inline dcontract{dim}(S1::Tensor{4, dim}, S2::SymmetricTensor{2, dim}) = dcontract(S1, convert(Tensor, S2))

@inline dcontract{dim}(S1::SymmetricTensor{2, dim}, S2::Tensor{2, dim}) = dcontract(promote(S1, S2)...)
@inline dcontract{dim}(S1::SymmetricTensor{2, dim}, S2::Tensor{4, dim}) = dcontract(convert(Tensor, S1), S2)

@inline dcontract{dim}(S1::Tensor{4, dim}, S2::SymmetricTensor{4, dim}) = dcontract(promote(S1, S2)...)
@inline dcontract{dim}(S1::SymmetricTensor{4, dim}, S2::Tensor{4, dim}) = dcontract(promote(S1, S2)...)

"""
```julia
otimes(::Vec, ::Vec)
otimes(::SecondOrderTensor, ::SecondOrderTensor)
```
Computes the open product between two tensors.
The symbol `⊗`, written `\\otimes`, is overloaded for tensor products.

**Example:**

```jldoctest
julia> A = rand(SymmetricTensor{2, 2});

julia> B = rand(SymmetricTensor{2, 2});

julia> A ⊗ B
2×2×2×2 ContMechTensors.SymmetricTensor{4,2,Float64,9}:
[:, :, 1, 1] =
 0.271839  0.352792
 0.352792  0.260518

[:, :, 2, 1] =
 0.469146  0.608857
 0.608857  0.449607

[:, :, 1, 2] =
 0.469146  0.608857
 0.608857  0.449607

[:, :, 2, 2] =
 0.504668  0.654957
 0.654957  0.48365
```
"""
@inline function otimes{dim}(S1::Tensor{2, dim}, S2::Tensor{2, dim})
    Tensor{4, dim}(tovector(S1) * tovector(S2)')
end

@inline function otimes{dim}(v1::Vec{dim}, v2::Vec{dim})
    Tensor{2, dim}(tovector(v1) * tovector(v2)')
end

const ⊗ = otimes

# Specialized methods for symmetric tensors
@inline function otimes{dim}(S1::SymmetricTensor{2, dim}, S2::SymmetricTensor{2, dim})
    SymmetricTensor{4, dim}(tovector(S1) * tovector(S2)')
end

# Promotion
@inline otimes{dim}(S1::SymmetricTensor{2, dim}, S2::Tensor{2, dim}) = otimes(promote(S1, S2)...)
@inline otimes{dim}(S1::Tensor{2, dim}, S2::SymmetricTensor{2, dim}) = otimes(promote(S1, S2)...)

"""
```julia
dot(::Vec, ::Vec)
dot(::Vec, ::SecondOrderTensor)
dot(::SecondOrderTensor, ::Vec)
dot(::SecondOrderTensor, ::SecondOrderTensor)
```
Computes the dot product (single contraction) between two tensors.
The symbol `⋅`, written `\\cdot`, is overloaded for single contraction.

**Example:**

```jldoctest
julia> A = rand(Tensor{2, 2})
2×2 ContMechTensors.Tensor{2,2,Float64,4}:
 0.590845  0.566237
 0.766797  0.460085

julia> B = rand(Tensor{1, 2})
2-element ContMechTensors.Tensor{1,2,Float64,2}:
 0.794026
 0.854147

julia> dot(A, B)
2-element ContMechTensors.Tensor{1,2,Float64,2}:
 0.952796
 1.00184

julia> A ⋅ B
2-element ContMechTensors.Tensor{1,2,Float64,2}:
 0.952796
 1.00184
```
"""
@inline Base.dot{dim}(v1::Vec{dim}, v2::Vec{dim}) = tovector(v1) ⋅ tovector(v2)

@inline function Base.dot{dim}(S1::Tensor{2, dim}, v2::Vec{dim})
    return Vec{dim}(tomatrix(S1) * tovector(v2))
end

@inline function Base.dot{dim}(v1::Vec{dim}, S2::Tensor{2, dim})
    return Vec{dim}(tomatrix(S2)' * tovector(v1))
end

@inline function Base.dot{dim}(S1::Tensor{2, dim}, S2::Tensor{2, dim})
    return Tensor{2, dim}(tomatrix(S1) * tomatrix(S2))
end

@inline function Base.dot{dim}(S1::SymmetricTensor{2, dim}, S2::SymmetricTensor{2, dim})
    S1_t = convert(Tensor{2, dim}, S1)
    S2_t = convert(Tensor{2, dim}, S2)
    return Tensor{2, dim}(tomatrix(S1_t) * tomatrix(S2_t))
end

@inline Base.dot{dim}(S1::SymmetricTensor{2, dim}, v2::Vec{dim}) = dot(convert(Tensor{2, dim}, S1), v2)

@inline Base.dot{dim}(v2::Vec{dim}, S1::SymmetricTensor{2, dim}) = dot(S1, v2)

# Promotion
Base.dot{dim}(S1::Tensor{2, dim}, S2::SymmetricTensor{2, dim}) = dot(promote(S1, S2)...)
Base.dot{dim}(S1::SymmetricTensor{2, dim}, S2::Tensor{2, dim}) = dot(promote(S1, S2)...)

"""
```julia
tdot(::SecondOrderTensor)
```
Computes the transpose-dot of a second order tensor with itself.
Returns a `SymmetricTensor`.

**Example:**

```jldoctest
julia> A = rand(Tensor{2,3})
3×3 ContMechTensors.Tensor{2,3,Float64,9}:
 0.590845  0.460085  0.200586
 0.766797  0.794026  0.298614
 0.566237  0.854147  0.246837

julia> tdot(A)
3×3 ContMechTensors.SymmetricTensor{2,3,Float64,6}:
 1.2577   1.36435   0.48726
 1.36435  1.57172   0.540229
 0.48726  0.540229  0.190334
```
"""
@generated function tdot{dim}(S1::Tensor{2, dim})
    idx(i,j) = compute_index(Tensor{2, dim}, i, j)
    ex = Expr(:tuple)
    for j in 1:dim, i in j:dim
        ex1, ex2 = Expr[], Expr[]
        for k in 1:dim
            push!(ex1, :(get_data(S1)[$(idx(k,i))]))
            push!(ex2, :(get_data(S1)[$(idx(k,j))]))
        end
        push!(ex.args, make_muladd_exp(ex1, ex2))
    end
    return quote
        $(Expr(:meta, :inline))
        @inbounds r = $ex
        SymmetricTensor{2, dim}(r)
    end
end

@inline tdot{dim}(S1::SymmetricTensor{2,dim}) = tdot(convert(Tensor{2,dim}, S1))

"""
```julia
dott(::SecondOrderTensor)
```
Computes the dot-transpose of a second order tensor with itself.
Returns a `SymmetricTensor`.

**Example:**

```jldoctest
julia> A = rand(Tensor{2,3})
3×3 ContMechTensors.Tensor{2,3,Float64,9}:
 0.590845  0.460085  0.200586
 0.766797  0.794026  0.298614
 0.566237  0.854147  0.246837

julia> dott(A)
3×3 ContMechTensors.SymmetricTensor{2,3,Float64,6}:
 0.601011  0.878275  0.777051
 0.878275  1.30763   1.18611
 0.777051  1.18611   1.11112
```
"""
@inline dott(S::SecondOrderTensor) = tdot(transpose(S))

"""
```julia
cross(::Vec, ::Vec)
```
Computes the cross product between two `Vec` vectors, returns a `Vec{3}`. For dimensions 1 and 2 the `Vec`'s
are expanded to 3D first. The infix operator `×` (written `\\times`) can also be used.

**Example:**

```jldoctest
julia> a = rand(Vec{3})
3-element ContMechTensors.Tensor{1,3,Float64,3}:
 0.590845
 0.766797
 0.566237

julia> b = rand(Vec{3})
3-element ContMechTensors.Tensor{1,3,Float64,3}:
 0.460085
 0.794026
 0.854147

julia> a × b
3-element ContMechTensors.Tensor{1,3,Float64,3}:
  0.20535
 -0.24415
  0.116354
```
"""
function Base.cross{T}(u::Vec{3, T}, v::Vec{3, T})
    @inbounds w = Vec{3}((u[2]*v[3] - u[3]*v[2], u[3]*v[1] - u[1]*v[3], u[1]*v[2] - u[2]*v[1]))
    return w
end
function Base.cross{T}(u::Vec{2, T}, v::Vec{2, T})
    @inbounds w = Vec{3}((zero(T), zero(T), u[1]*v[2] - u[2]*v[1]))
    return w
end
function Base.cross{T}(::Vec{1, T}, ::Vec{1, T})
    return zero(Vec{3,T})
end


struct Input{T}
    year::Vector{Int}
    month::Vector{Int}
    day::Vector{Int}
    hour::Vector{Int}
    SW::Vector{T}
    LW::Vector{T}
    Sf::Vector{T}
    Rf::Vector{T}
    Ta::Vector{T}
    RH::Vector{T}
    Ua::Vector{T}
    Ps::Vector{T}
end

Base.@kwdef struct Constants{T}

    cp::T = 1005          # Specific heat capacity of dry air (J/K/kg)
    eps::T = 0.622       # Ratio of molecular weights of water and dry air    BEFORE eps
    e0::T = 611.213       # Saturation vapour pressure at Tm (Pa)
    g::T = 9.81           # Acceleration due to gravity (m/s^2)
    hcap_ice::T = 2100.   # Specific heat capacity of ice (J/K/kg)
    hcap_wat::T = 4180.   # Specific heat capacity of water (J/K/kg)
    hcon_air::T = 0.025   # Thermal conductivity of air (W/m/K)
    hcon_clay::T = 1.16   # Thermal conductivity of clay (W/m/K)
    hcon_ice::T = 2.24    # Thermal conducivity of ice (W/m/K)
    hcon_sand::T = 1.57   # Thermal conductivity of sand (W/m/K)
    hcon_wat::T = 0.56    # Thermal conductivity of water (W/m/K)
    Lc::T = 2.501e6       # Latent heat of condensation (J/kg)
    Lf::T = 0.334e6       # Latent heat of fusion (J/kg)
    Ls::T = Lc + Lf       # Latent heat of sublimation (J/kg)
    R::T = 8.3145         # Universal gas constant (J/K/mol)
    Rgas::T = 287         # Gas constant for dry air (J/K/kg)
    Rwat::T = 462         # Gas constant for water vapour (J/K/kg)
    rho_ice::T = 917.     # Density of ice (kg/m^3)
    rho_wat::T = 1000.    # Density of water (kg/m^3)
    sb::T = 5.67e-8       # Stefan-Boltzmann constant (W/m^2/K^4)
    Tm::T = 273.15        # Melting point (K)
    vkman::T = 0.4        # Von Karman constant

end

Base.@kwdef mutable struct EBM{T}

    ### Settings ##############################################################

    # Maximum number of snow layers
    Nsmax::Int = 3

    # Number of soil layers
    Nsoil::Int = 4

    # Time step (s)
    dt::T = 3600

    # Minimum snow layer thicknesses (m)
    Dzsnow::Vector{T} = [0.1, 0.2, 0.4]

    # Soil layer thicknesses (m)
    Dzsoil::Vector{T} = [0.1, 0.2, 0.4, 0.8]

    # Temperature measurement height (m)
    zT::T = 2

    # Wind measurement height (m)
    zU::T = 10

    # Subtract snow depth from measurement height
    zvar::Bool = true

    ### Snow parameters #######################################################

    # Maximum albedo for fresh snow
    asmx::T = 0.8

    # Minimum albedo for melting snow
    asmn::T = 0.5

    # Stability slope parameter
    bstb::T = 5

    # Snow thermal conductivity exponent
    bthr::T = 2

    # Snow cover fraction depth scale (m)
    hfsn::T = 0.1

    # Fixed thermal conductivity of snow (W/m/K)
    kfix::T = 0.24

    # Fixed snow density (kg/m^3)
    rho0::T = 300

    # Fresh snow density (kg/m^3)
    rhof::T = 100

    # Maximum density for cold snow (kg/m^3)
    rcld::T = 300

    # Maximum density for melting snow (kg/m^3)
    rmlt::T = 500

    # Snowfall to refresh albedo (kg/m^2)
    Salb::T = 10

    # Albedo decay temperature threshold (C)
    Talb::T = -2

    # Cold snow albedo decay timescale (h)
    tcld::T = 1000

    # Melting snow albedo decay timescale (h)
    tmlt::T = 100

    # Snow compaction time scale (h)
    trho::T = 200

    # Irreducible liquid water content of snow
    Wirr::T = 0.03

    # Snow roughness length (m)
    z0sn::T = 0.01

    ### Surface parameters ####################################################

    # Snow-free ground albedo
    alb0::T = 0.2

    # Surface conductance for saturated soil (m/s)
    gsat::T = 0.01

    # Snow-free roughness length (m)
    z0sf::T = 0.1

    ### Soil properties #######################################################

    # Soil clay fraction
    fcly::T = 0.3

    # Soil sand fraction
    fsnd::T = 0.6

    # Clapp-Hornberger exponent
    b::T = 7.630000000000001

    # Volumetric heat capacity of dry soil (J/K/m^3)
    hcap_soil::T = 2.2993333333333335e6

    # Saturated soil water pressure (m)
    sathh::T = 0.10789467222298288

    # Volumetric soil moisture concentration at saturation
    Vsat::T = 0.4087

    # Volumetric soil moisture concentration at critical point
    Vcrit::T = 0.2603859120641063

    # Thermal conductivity of dry soil (W/m/K)
    hcon_soil::T = 0.2740041303112452

    ### State variables #######################################################

    # Snow albedo
    albs::T = 0.8

    # Snow layer thicknesses (m)
    Ds::Vector{T} = fill(0, Nsmax)

    # Number of snow layers
    Nsnow::Int = 0

    # Ice content of snow layers (kg/m^2)
    Sice::Vector{T} = fill(0, Nsmax)

    # Liquid content of snow layers (kg/m^2)
    Sliq::Vector{T} = fill(0, Nsmax)

    # Volumetric moisture content of soil layers
    theta::Vector{T} = fill(0.5 * Vsat, Nsoil)

    # Snow layer temperatures (K)
    Tsnow::Vector{T} = fill(273.15, Nsmax)

    # Soil layer temperatures (K)
    Tsoil::Vector{T} = fill(285, Nsoil)

    # Surface skin temperature (K)
    Tsurf::T = Tsoil[1]

    ### Configurations ########################################################

    # Snow albedo model
    am::Int = 0

    # Snow conductivity model
    cm::Int = 0

    # Snow density model
    dm::Int = 0

    # Surface exchange model
    em::Int = 0

    # Snow hydraulics model
    hm::Int = 0

    ### Internal variables ####################################################

    # Thermal conductivity of snow (W/m/K)
    ksnow::Vector{T} = fill(0, Nsmax)

    # Thermal conductivity of soil (W/m/K)
    ksoil::Vector{T} = fill(0, Nsoil)

    # Areal heat capacity of soil (J/K/m^2)
    csoil::Vector{T} = fill(0, Nsoil)

    # Surface moisture conductance (m/s)
    gs::T = 0

    # Effective albedo
    alb::T = 0

    # Fresh snow density (kg/m^3)
    rfs::T = 0

    # Snow cover fraction
    fsnow::T = 0

    # Transfer coefficient for heat and moisture
    CH::T = 0

    # Surface roughness length (m)
    z0::T = 0

    # Surface thermal conductivity (W/m/K)
    ksurf::T = 0

    # Surface layer temperature (K)
    Ts1::T = 0

    # Surface layer thickness (m)
    Dz1::T = 0

    # Snow sublimation rate (kg/m^2/s)
    Esnow::T = 0

    # Heat flux into surface (W/m^2)
    Gsurf::T = 0

    # Sensible heat flux (W/m^2)
    Hsurf::T = 0

    # Latent heat flux (W/m^2)
    LEsrf::T = 0

    # Surface melt rate (kg/m^2/s)
    Melt::T = 0

    # Net radiation (W/m^2)
    Rnet::T = 0

    # Heat flux into soil (W/m^2)
    Gsoil::T = 0

    # Temporary soil variables
    soil_a::Vector{T} = fill(0, Nsoil)
    soil_b::Vector{T} = fill(0, Nsoil)
    soil_c::Vector{T} = fill(0, Nsoil)
    soil_Gs::Vector{T} = fill(0, Nsoil)
    soil_rhs::Vector{T} = fill(0, Nsoil)
    soil_dTs::Vector{T} = fill(0, Nsoil)
    soil_gamma::Vector{T} = fill(0, Nsoil)

    # Temporary snow variables
    snow_csnow::Vector{T} = fill(0, Nsmax)
    snow_E::Vector{T} = fill(0, Nsmax)
    snow_U::Vector{T} = fill(0, Nsmax)
    snow_Gs::Vector{T} = fill(0, Nsmax)
    snow_dTs::Vector{T} = fill(0, Nsmax)
    snow_a::Vector{T} = fill(0, Nsmax)
    snow_b::Vector{T} = fill(0, Nsmax)
    snow_c::Vector{T} = fill(0, Nsmax)
    snow_rhs::Vector{T} = fill(0, Nsmax)
    snow_gamma::Vector{T} = fill(0, Nsmax)

end


Ψ = qnmfunction(s=2,l=2,m=2,n=2,a=0.1)
Ψ2 = qnmfunction(s=-1,l=2,m=2,n=2,a=0.1)

Ψ2 = qnmfunction(s=2,l=2,m=-2,n=3,a=0.7)
Ψ2

using Plots
using LaTeXStrings
r = (Ψ.R.r₊+0.2):0.001:10
rʼ = 8.0:0.1:10
plotlyjs()
using Statistics
let lw=1, height = 3
Ψ1 = qnmfunction(s=-1,l=2,m=2,n=2,a=0.1)
Ψ2 = qnmfunction(s=2,l=2,m=-2,n=0,a=0.2)
rʼ = 2.0:0.1:10
wave1 = real.(Ψ1.(rʼ,0.01,0.01,2.0))
wave2 = real.(Ψ2.(rʼ,0.01,0.01,2.0))
wh1 = mean(abs.(wave1))
wh2 = mean(abs.(wave2))
#wh = height*max(wave1height,wave2height)
p = plot(rʼ,wave1./wh1,
        ylim = [-height,height],
        lab = L"{\ }_{%$(Ψ1.s)}\Psi_{%$(Ψ1.l),%$(Ψ1.m)}",
        linewidth = lw)

plot!(p,rʼ,wave2./wh2,
        ylim = [-height,height],
        lab = L"{\ }_{%$(Ψ2.s)}\Psi_{%$(Ψ2.l),%$(Ψ2.m)}",
        linewidth = lw)

end

gr()
Ψm2 = qnmfunction(s=-2,l=2,m=2,n=0,a=0.1)
Ψm1 = qnmfunction(s=-2,l=2,m=1,n=0,a=0.1)
Ψ0 = qnmfunction(s=-2,l=2,m=0,n=0,a=0.1)
Ψp1 = qnmfunction(s=-2,l=2,m=-1,n=0,a=0.1)
Ψp2 = qnmfunction(s=-2,l=2,m=-2,n=0,a=0.1)
ΨArray = [Ψm2,Ψm1,Ψ0,Ψp1,Ψp2]
rʼ = 2.0:0.1:10
lw=1; height = 3;
anim = @animate for t ∈ 0:0.2:3
        N = length(ΨArray)
        waves = [];
        whs = [];
        for i in 1:N
            thewave = real.(ΨArray[i].R.(rʼ)*exp(-ΨArray[i].ω*im*t))
            push!(waves,thewave)
            push!(whs,mean(abs.(thewave)))
        end

    wh = height*maximum(whs)
    p = plot(rʼ,waves[1],
            ylim = [-200,200],
            lab = L"{\ }_{%$(ΨArray[1].s)}R_{%$(ΨArray[1].l),%$(ΨArray[1].m)}",
            linewidth = lw)
        for i in 2:N
            plot!(p,rʼ,waves[i],
                ylim = [-200,200],
                lab = L"{\ }_{%$(ΨArray[i].s)}R_{%$(ΨArray[i].l),%$(ΨArray[i].m)}",
                linewidth = lw)
        end
        xlabel!(p,L"r (\textrm{in units of M})")
        ylabel!(p,L"{\ }_{s}R_{lm}")
        title!(p, "Quasinormal Mode Radial Function time dependence \n Does blow up at infinity")

end
display(anim)
gif(anim, "anim_fps15.gif", fps = 30)

x = 0:0.1:100
ω = 0.4 - im*0.5
anim2 = @animate for t ∈ 0:1:100
        plot(x, map(r-> real(exp(-im*ω*(t-r))),x), ylim = [-1000,1000])
end
gif(anim2, "anim_fps15.gif", fps = 30)

Ψ2(Ψ2.R.r₊+0.000001)

function factorial2(n)
    if n > 20
        return n
    return factorial(n)
end

function SterlingsApprox(n)
    return sqrt(2*pi*n)*(n/e)^n
end

#=
using PyCall
qnm = pyimport("qnm")
grav_220 = qnm.modes_cache(s=-2,l=2,m=2,n=2)
omega, A, C = grav_220(a=0.1)
println(omega)
Ψ1 = qnmfunction(s=-2,l=2,m=2,n=2,a=0.1)
println(Ψ1.ω)
grav_220 = qnm.modes_cache(s=-2,l=2,m=-2,n=2)
omega, A, C = grav_220(a=0.1)
println(omega)
Ψ2 = qnmfunction(s=-2,l=2,m=-2,n=2,a=0.1)
println(Ψ2.ω)
=#


using OrbitalTrajectories

using Test

@testset "SpiceUtils" begin
    @test OrbitalTrajectories.SpiceUtils.mass_fraction("jupiter", "europa") ≈ 2.528e-5 atol=1e-5
    @test OrbitalTrajectories.SpiceUtils.mass_fraction("mars", "phobos") ≈ 1.0659e16 / (1.0659e16 + 6.4171e23) atol=1e-5
end


using ITensors,
      ITensorsGPU,
      LinearAlgebra, # For tr()
      Combinatorics, # For permutations()
      CuArrays,
      Test

      # gpu tests!
@testset "cuITensor, Dense{$SType} storage" for SType ∈ (Float64, ComplexF64)
  mi,mj,mk,ml,ma = 2,3,4,5,6,7
  i = Index(mi,"i")
  j = Index(mj,"j")
  k = Index(mk,"k")
  l = Index(ml,"l")
  a = Index(ma,"a")
  @testset "Constructor" begin
      A = cuITensor(one(SType),i,j,k)
      @test collect(CuArray(A, i, j, k)) == ones(SType, dim(i), dim(j), dim(k))
      A = randomCuITensor(IndexSet(i,j,k))
      @test inds(A) == IndexSet(i, j, k)
      @test ITensorsGPU.store(A) isa ITensorsGPU.CuDense
      Aarr = rand(SType, dim(i)*dim(j)*dim(k))
      @test collect(ITensor(Aarr, i, j, k)) == collect(cuITensor(Aarr, i, j, k))
      @test cuITensor(SType, i, j, k) isa ITensor{3}
      @test store(cuITensor(SType, i, j, k)) isa ITensorsGPU.CuDense{SType}
      @test vec(collect(CuArray(ITensor(Aarr, i,j,k), i, j, k))) == Aarr
  end
  @testset "Test permute(cuITensor,Index...)" begin
    CA     = randomCuITensor(SType,i,k,j)
    permCA = permute(CA,k,j,i)
    permA  = collect(permCA)
    @test k==inds(permA)[1]
    @test j==inds(permA)[2]
    @test i==inds(permA)[3]
    A = collect(CA)
    for ii ∈ 1:dim(i), jj ∈ 1:dim(j), kk ∈ 1:dim(k)
      @test A[k(kk),i(ii),j(jj)]==permA[i(ii),j(jj),k(kk)]
    end
  end
  #=@testset "Test scalar(cuITensor)" begin
    x = SType(34)
    A = randomCuITensor(a)
    @test x==scalar(A)
  end=#
  @testset "Test CuVector(cuITensor)" begin
      v = CuVector(ones(SType, dim(a)))
      A = cuITensor(v, a)
      @test v==CuVector(A)
  end
  @testset "Test CuMatrix(cuITensor)" begin
      v = CuMatrix(ones(SType, dim(a), dim(l)))
      A = cuITensor(vec(v), a, l)
      @test v==CuMatrix(A, a, l)
      A = cuITensor(vec(v), a, l)
      @test v==CuMatrix(A)
      A = cuITensor(vec(v), a, l)
      @test v==CuArray(A, a, l)
      @test v==CuArray(A)
  end
  @testset "Test norm(cuITensor)" begin
    A = randomCuITensor(SType,i,j,k)
    B = dag(A)*A
    @test norm(A)≈sqrt(real(scalar(B)))
  end
  @testset "Test complex(cuITensor)" begin
    A  = randomCuITensor(SType,i,j,k)
    cA = complex(A)
    @test complex.(CuArray(A)) == CuArray(cA)
  end
  #@testset "Test exp(cuITensor)" begin
  #  A  = randomCuITensor(SType,i,i')
  #  @test CuArray(exp(A,i,i')) ≈ exp(CuArray(A))
  #end
  @testset "Test add cuITensors" begin
    dA = randomCuITensor(SType,i,j,k)
    dB = randomCuITensor(SType,k,i,j)
    A = collect(dA)
    B = collect(dB)
    C = collect(dA+dB)
    @test CuArray(permute(C,i,j,k))==CuArray(permute(A,i,j,k))+CuArray(permute(B,i,j,k))
    for ii ∈ 1:dim(i), jj ∈ 1:dim(j), kk ∈ 1:dim(k)
      @test C[i(ii),j(jj),k(kk)]==A[j(jj),i(ii),k(kk)]+B[i(ii),k(kk),j(jj)]
    end
  end

  @testset "Test factorizations of a cuITensor" begin

    A = randomCuITensor(SType,i,j,k,l)

    @testset "Test SVD of a cuITensor" begin
      U,S,V = svd(A,(j,l))
      u = commonind(U,S)
      v = commonind(S,V)
      @test collect(A)≈collect(U*S*dag(V))
      @test collect(U*dag(prime(U,u)))≈δ(SType,u,u') rtol=1e-14
      @test collect(V*dag(prime(V,v)))≈δ(SType,v,v') rtol=1e-14
    end

    #=@testset "Test SVD truncation" begin
        M = randn(4,4)
        (U,s,V) = svd(M)
        ii = Index(4)
        jj = Index(4)
        S = Diagonal(s)
        T = cuITensor(vec(CuArray(U*S*V')),IndexSet(ii,jj))
        (U,S,V) = svd(T,ii;maxm=2)
        @test norm(U*S*V-T)≈sqrt(s[3]^2+s[4]^2)
    end=#

    @testset "Test QR decomposition of a cuITensor" begin
      Q,R = qr(A,(i,l))
      q = commonind(Q,R)
      @test collect(A)≈collect(Q*R)
      @test collect(Q*dag(prime(Q,q)))≈δ(SType,q,q') atol=1e-14
    end

    @testset "Test polar decomposition of a cuITensor" begin
      U,P = polar(A,(k,l))
      @test collect(A)≈collect(U*P)
    end

  end # End ITensor factorization testset
end # End Dense storage test


using Unitful, UnitfulRecipes, UnitfulAstro
using Dierckx

#
# The file containing the various differential equation systems to solve.
#

#--------------------------
# Definition of equations
#--------------------------

"The differential equation system derived by Draine 2011"
function d2011system!(dv, v, p, y)
    u, ϕ, τ = v
    γ, β = p

    dv[1] = du = -1 - γ*(β*exp(-τ) + ϕ)*u/y^2
    dv[2] = dϕ = -(y/u)^2 - γ * ϕ/u
    dv[3] = dτ = γ/u
end

"""
    draine2011model(γ, β; u0=454, ystart=10, yend=100)

Solve the Draine (2011) model with model parameters γ and β.
This formulates the ODEProblem and solves it

"""
function draine2011model(γ, β; u0=454.0, ystart=10.0, yend=100.0)

    starting_point = [u0, 1.0, 0.0]
    yspan = [ystart, yend]

    p = [γ, β]

    # The system is only defined for ϕ>=0 so we interrupt on ϕ=u[2]=0
    condition(u, y, integrator) = u[2]
    affect!(integrator) = terminate!(integrator)
    cb = ContinuousCallback(condition, affect!);
    prob = ODEProblem(d2011system!, starting_point, yspan, p)

    sol = solve(prob, Tsit5(), callback=cb);

    result = Draine2011(sol, Draine2011Parameters(γ, β), u0,
                        starting_point[2], starting_point[3],
                        ystart, ymax(sol))

    return result

end

draine2011model(d::NamedTuple; kwargs...) = draine2011model(d.γ, d.β; u0=d.u0, kwargs...)

function d2011wgravity!(dv, v, h, p, y)
    u, ϕ, τ = v
    γ, β, η, ymin = p # Ymin is indirectly a parameter of the model

    rint(x) = x^2/h(p, x)[1]
    Ry, err = quadgk(rint, ymin, y)

    dv[1] = du = u^2*η*Ry/y^2 -1 - γ*(β*exp(-τ) + ϕ)*u/y^2
    dv[2] = dϕ = -(y/u)^2 - γ * ϕ/u
    dv[3] = dτ = γ/u
end

function d2011wgravitymodel(γ, β, η; u0=454, ystart=10, yend=100)

    starting_point = [u0, 1, 0]
    yspan = [ystart, yend]

    p = [γ, β, η, ystart]

    # The system is only defined for ϕ>=0 so we interrupt on ϕ=u[2]=0
    condition(u, y, integrator) = u[2]
    affect!(integrator) = terminate!(integrator)
    cb = ContinuousCallback(condition, affect!);

    # This specifies the original starting point. I think it is probably
    # necessary to rethink this  but haven't done this yet
    history(p, t) = [u0, 1, 0]
    prob = DDEProblem(d2011wgravity!, starting_point, history, yspan, p)

    # The choice for stepper is following https://diffeq.sciml.ai/stable/solvers/dde_solve/
    # see also https://github.com/SciML/DifferentialEquations.jl/issues/383
    sol = solve(prob, MethodOfSteps(RK4()), callback=cb);

    return sol

end

#
# Handling of solutions themselves
#

ymax(solution) = maximum(solution.t)
ymin(solution) = minimum(solution.t)

# Return the functions of the differential equation
tau(y, model::AbstractDustyHII) = model.solution(y; idxs=3)
phi(y, model::AbstractDustyHII) = model.solution(y; idxs=2)
ufunc(y, model::AbstractDustyHII) = model.solution(y; idxs=1)

#
# To-from initial conditions.
#
# In this case we have two options:
#
#    1. The user might give physical parameters of a model and want
#       the parameters to use for initial conditions and model parameters γ and β.
#
#    2. The user might have a solution for a given γ and β and wants physical
#       parameters that are consistent with this. This is the more complex case
#       because of the various extra parameters that might have to be passed.
#

"""
    physical_to_model(; Teff=45000u"K", Tgas=1e4u"K", ninner=1e3u"1/cm^3")

Calculates β, Q0, hν_i for a star with the given Teff and calculates γ and
η for the other properties.

"""
function physical_to_model(;Teff=45000u"K", Tgas=1e4u"K", ninner=1e3u"1/cm^3",
                           σ=1e-21*u"1/cm^2")

    # Calculate first a Black-body spectrum - I fix wavelength ranges to
    # be reasonable for a hot star.
    lrange = collect(LinRange(1., 6e4, 10000))
    flux = [PlanckAstro(l*u"angstrom", Teff) for l in lrange];
    itp = interpolate_spectrum(lrange, flux)
    resStar = specQuantities(lrange, itp)

    # Then calculate γ
    γ = gamma(;T=Tgas, Q0=resStar.Q0, hnui=resStar.hnui, σ=σ)

    #  η
    η = upreferred(Eg(Tgas; Q0=resStar.Q0, hnui=resStar.hnui, σ=σ)/Et(Tgas))

    # u0
    u0 = ustrip(uval(ninner; T=Tgas, Q0=resStar.Q0, hnui=resStar.hnui))

    return (Ln=resStar.Ln, Li=resStar.Li, hnui=resStar.hnui,
            Q0=resStar.Q0, β=resStar.β, γ=γ, η=η, u0=u0, Teff=Teff,
            Tgas=Tgas, ninner=ninner, σ=σ)

end

"""
    model_to_physical(m)

Takes a model which has been calculated in unitless form and converts quantities
to physical units.

NOT YET FUNCTIONAL
"""
function model_to_physical(m, T, Q0, hnui, σ)
    return Draine2011Physical(m, T, Q0, hnui, σ)

end


function parserule(rule)
    ss = split(rule, " bags contain ")
    rule = ss[1]
    edges = Vector{typeof(rule)}()
    rest = split(ss[2], ", ")
    rx = r"([0-9]+) ([a-z]+ [a-z]+)"
    for s in rest
        m = match(rx, s)
        if m == nothing
            continue
        end
        num = parse(Int, m.captures[1])
        append!(edges, repeat([m.captures[2]], num))
    end
    (rule, edges)
end

function parserules(rules)
    ss = split(rules, "\n", keepempty=false)
    map(parserule, ss)
end

#=
parserule("light red bags contain 1 bright white bag, 2 muted yellow bags.")
parserule("bright white bags contain 1 shiny gold bag.")
parserule("faded blue bags contain no other bags.")

rules = """
light red bags contain 1 bright white bag, 2 muted yellow bags.
dark orange bags contain 3 bright white bags, 4 muted yellow bags.
bright white bags contain 1 shiny gold bag.
muted yellow bags contain 2 shiny gold bags, 9 faded blue bags.
shiny gold bags contain 1 dark olive bag, 2 vibrant plum bags.
dark olive bags contain 3 faded blue bags, 4 dotted black bags.
vibrant plum bags contain 5 faded blue bags, 6 dotted black bags.
faded blue bags contain no other bags.
dotted black bags contain no other bags.
"""
parsed_rules = parserules(rules)
=#

function build_graph(parsed_rules::Vector{Tuple{SubString{String}, Vector{SubString{String}}}})
    m = length(parsed_rules)
    # Rows indicate what bags 'this' bag may contain. Columns indicate what bags may
    # contain 'this' bag.
    M = zeros(Int64, (m, m))
    idx_tracker = Dict{AbstractString, Int64}()
    # Just build index tracker first to make updating matrix easy.
    for (i, (bag, _)) in enumerate(parsed_rules)
        idx_tracker[bag] = i
    end
    # Build up matrix.
    for (i, (bag, contains_bags)) in enumerate(parsed_rules)
        for contains_bag in contains_bags
            j = idx_tracker[contains_bag]
            M[i, j] += 1
        end
    end
    (idx_tracker, M)
end

#=
(idxs, M) = build_graph(parsed_rules)
idxs
M
=#

function dfs_walk(f, M, col_idx)
    for (i, val) in enumerate(M[:, col_idx])
        if val == 0
            continue
        end
        f(i, col_idx, val)
        dfs_walk(f, M, i)
    end
end

function coalesce_bag_vals(M, row_idx)
    acc = sum(M[row_idx, :])
    if acc == 0
        return 0
    end
    for (j, val) in enumerate(M[row_idx, :])
        if val == 0
            continue
        end
        acc += val * coalesce_bag_vals(M, j)
    end
    acc
end

#=
dfs_walk((i, j, val)->println("walked $(i) $(j) $(val)"), M, 5)
coalesce_bag_vals(M, 5)
=#

function num_bags_contain(M, idxs, bagstr)
    if !haskey(idxs, bagstr)
        error("No index for bag")
    end
    idx = get(idxs, bagstr, 0)
    walked_bags = Vector{Int64}(undef, 0)
    dfs_walk((i, _, _) -> push!(walked_bags, i), M, idx)
    length(unique(walked_bags))
end

#=
num_bags_contain(M, idxs, "shiny gold")
bag_vals = num_bags_contain_val(M, idxs, "shiny gold")
=#

function day07part1(s)
    parsed_rules = parserules(s)
    (idxs, M) = build_graph(parsed_rules)
    num_bags_contain(M, idxs, "shiny gold")
end

function day07part2(s)
    parsed_rules = parserules(s)
    (idxs, M) = build_graph(parsed_rules)
    coalesce_bag_vals(M, get(idxs, "shiny gold", 0))
end

function day07(file)
    s = read(file, String)
    (day07part1(s), day07part2(s))
end

#=
day07("input.txt")
=#


"""
Tests for R2BP dynamics.
"""
module R2BPTests

    using ModelingToolkit, AstrodynamicalModels, Test

    @testset "R2BP Model Constructors" begin
        model = R2BP(; stm=false, structural_simplify=false)
        model = R2BP(; stm=true, structural_simplify=true)
        @test true
    end

    @testset "R2BP Model Calculations" begin
        vectorfield = R2BPFunction()

        r = [-11e3, 0, 5e3]
        v = [0.0, 0.0, 0.0]
        μ = 398600.4354360959

        @test isapprox(
                vectorfield(vcat(r,v), [μ], NaN),
                [0.0, 0.0, 0.0, 0.00248543, -0.0, -0.00112974];
                atol = 1e-8
        )
    end

end


# define constants
using PyPlot

const g = 10.
const H = 10.
const N = 500
const Nt = 3000
const cfl = 0.9
const rho = 1.
const ν = 0.006
const F0 = 500.
const Lx = 1.
const dx = Lx/N
const dxinv = 1./dx
const cph = sqrt(g*H)
const dt = cfl * dx / cph
const RKα = [1/6.,1/3.,1/3.,1/6.]
const RKβ = [0.5,0.5,1.]

# grid
const x_h = dx/2:dx:Lx
const x_u = dx:dx:Lx        # no u-point at x=0 but at x=L (periodicity)

F(x,t) = -F0*sin.(8π*x/Lx + 200*t).*sin.(π*x/Lx).^2/rho

# OPERATORS
function Gux(res,u)
    res[1] = dxinv*(u[1]-u[end])
    res[2:end] = dxinv*(u[2:end]-u[1:end-1])
    return res
end

function GTx(res,η)
    res[1:end-1] = dxinv*(η[2:end]-η[1:end-1])
    res[end] = dxinv*(η[1]-η[end])
    return res
end

function ITu(res,η)
    res[1:end-1] = .5*(η[1:end-1] + η[2:end])
    res[end] = .5*(η[1]+η[end])
    return res
end

function IuT(res,u)
    res[1] = .5*(u[1] + u[end])
    res[2:end] = .5*(u[1:end-1]+u[2:end])
    return res
end

# initial conditions
η0 = zeros(N)
u0 = zeros(N)

function rhs(du,dη,h_u,u_h,dudx,u,η,t)
    h_u = ITu(h_u,η+H)
    u_h = IuT(u_h,u.^2)
    dudx = Gux(dudx,u)

    du = -GTx(du,.5*u_h + g*η - ν*dudx)
    du += F(x_h,t)./h_u

    dη = -Gux(dη,u.*h_u)

    return du,dη
end

function time_integration(Nt,u,η)

    # pre-allocate memory
    u0,η0 = zeros(u),zeros(η)
    u1,η1 = zeros(u),zeros(η)
    du,dη = zeros(u),zeros(η)
    h_u = zeros(u)
    u_h = zeros(η)
    dudx = zeros(η)
    t = 0.

    # for output
    u_out = zeros(Nt+1,length(u))
    η_out = zeros(Nt+1,length(η))

    # store the initial conditions
    u_out[1,:] = u
    η_out[1,:] = η

    for i = 1:Nt
        u1[:] = u
        η1[:] = η

        for RKi = 1:4
            du,dη = rhs(du,dη,h_u,u_h,dudx,u1,η1,t)

            if RKi < 4 # RHS update for the next RK-step
                u1 = u + RKβ[RKi]*dt*du
                η1 = η + RKβ[RKi]*dt*dη
            end

            # Summing all the RHS on the go
            u0 += RKα[RKi]*dt*du
            η0 += RKα[RKi]*dt*dη

        end

        u[:] = u0
        η[:] = η0
        t += dt

        # # output
        u_out[i+1,:] = u
        η_out[i+1,:] = η

    end
    u_out,η_out
end

time_integration(1,u0,η0)
@time u,η = time_integration(Nt,u0,η0)

## PLOTTING
fig,ax = subplots()

l1, = ax[:plot](x_h,η[1,:])
ax[:set_ylim](-2,2)
ax[:set_xlabel]("x")
ax[:set_ylabel]("η")
tight_layout()

for it = 2:5:Nt+1
    pause(0.00001)
    l1[:set_data](x_h,η[it,:])
end

pause(2)
close(fig)


"""
    likelihood(events::Array{<:Number, 1}, background::Float64, kappa::Float64, kernel::Distributions.Distribution, maxT::Number)
Calculate the log likelihood of a Hawkes process for given parameters.

# Arguments

* `events` Vector of events to calculate the likelihood for.
* `background` Background rate.
* `kappa` Kappa parameter.
* `kernel` Function or distribution of the kernel.
* `maxT` Maximum time of the process.

# Notes

* The kernel function must be a proper probability distribution.

# Examples

"""
function likelihood(events::Array{<:Number, 1}, background::Float64, kappa::Float64, kernel::Distributions.Distribution, maxT::Number)
    kernel_f, kernel_cdf = generate_kernel_functions(kernel)
    likelihood(events, background, kappa, kernel_f, kernel_cdf, maxT)
end

function likelihood(events::Array{<:Number, 1}, background::Float64, kappa::Float64, kernel::Function, kernel_cdf::Function, maxT::Number)

    logIntensitySum = sum(log.(intensity(events, events, background, kappa, kernel)))
    backgroundIntegral = background * maxT
    kernelIntegral = sum(kernel_cdf(maxT .- events))

    logIntensitySum - backgroundIntegral - kappa*kernelIntegral
end


module Euclid
import ProjectiveGeometricAlgebra3d as PGA
using .PGA
using .PGA: sandwich
using .PGA: MultiVector3, MultiVector1, MultiVector2, MultiVectorEven, ∨
using .PGA: normalize, scalartype, normalize_ideal
using .PGA: random
using ArgCheck

export Point, Plane, Direction, Line
export meet, orthogonal_through, join, distance, project
export Motor, translator, rotator
export random

function pga end

struct Point{T}
    data::MultiVector3{T}
    function Point(data::MultiVector3)
        data1 = normalize(data)
        T = scalartype(typeof(data1))
        return new{T}(data1)
    end
end
Point(pt::Point) = pt
function Point(itr)
    @argcheck length(itr) == 3
    x,y,z = itr
    Point(x,y,z)
end
function Point(x,y,z)
    Point(MultiVector3(e023=-x, e013=y, e012=-z, e123=oneunit(x)))
end

struct Direction{T}
    data::MultiVector3{T}
    function Direction(data::MultiVector3)
        data1 = normalize_ideal(data)
        T = scalartype(typeof(data1))
        return new{T}(data1)
    end
end
Direction(o::Direction) = o
function Direction(itr)
    @argcheck length(itr) == 3
    x,y,z = itr
    Direction(x,y,z)
end
function Direction(x,y,z)
    data::MultiVector3 = MultiVector3(e023=-x, e013=y, e012=-z)
    Direction(data)
end

struct Line{T}
    data::MultiVector2{T}
    function Line(data::MultiVector2)
        data1 = normalize(data)
        T = scalartype(typeof(data1))
        return new{T}(data1)
    end
end
Line(o::Line) = o
function Line(;point=nothing, direction=nothing)
    has_point = point !== nothing
    has_direction = direction !== nothing
    if has_point && has_direction
        Line_point_direction(point, direction)
    else
        throw(ArgumentError("Cannot construct Line from arguments"))
    end
end
function Line_point_direction(point, direction)
    data::MultiVector2 = pga(Point(point)) ∨ pga(Direction(direction))
    Line(data)
end

struct Plane{T}
    data::MultiVector1{T}
    function Plane(data::MultiVector1)
        data1 = normalize(data)
        T = scalartype(typeof(data1))
        return new{T}(data1)
    end
end
Plane(o::Plane) = o
function Plane(;point, directions)
    Plane_point_directions(point, directions)
end
function Plane_point_directions(point, directions)
    @argcheck length(directions) == 2
    dir1, dir2 = directions
    pt::MultiVector3 = pga(Point(point))
    d1::MultiVector3 = pga(Direction(dir1))
    d2::MultiVector3 = pga(Direction(dir2))
    data::MultiVector1 = pt ∨ d1 ∨ d2
    Plane(data)
end

################################################################################
##### Motor
################################################################################
struct Motor{T}
    data::MultiVectorEven{T}
end
Motor(m::Motor) = m

function rotator(;axis, angle)
    l = Line(axis)
    s,c = sincos(1//2*angle)
    data = c + s*pga(l)
    return Motor(data)
end

function translator end

(m::Motor)(o::Point) = Point(sandwich(pga(o), pga(m)))
(m::Motor)(o::Direction) = Direction(sandwich(pga(o), pga(m)))
(m::Motor)(o::Line)  = Line(sandwich(pga(o), pga(m)))
(m::Motor)(o::Plane) = Plane(sandwich(pga(o), pga(m)))
(m::Motor)(o::Motor) = Motor(sandwich(pga(o), pga(m)))

Base.:(∘)(m2::Motor, m1::Motor) = Motor(pga(m1) * pga(m2))

pgatype(::Type{Point{T}}     ) where {T} = MultiVector3{T}
pgatype(::Type{Direction{T}} ) where {T} = MultiVector3{T}
pgatype(::Type{Line{T}}      ) where {T} = MultiVector2{T}
pgatype(::Type{Plane{T}}     ) where {T} = MultiVector1{T}
pgatype(::Type{Motor{T}}     ) where {T} = MultiVectorEven{T}

for Geo in [:Point, :Direction, :Line, :Plane, :Motor]
    @eval function pga(x::$Geo)
        x.data
    end
    @eval function PGA.random(rng::PGA.AbstractRNG, G::Type{<:$Geo})
        $Geo(PGA.random(rng, pgatype(G)))::G
    end
    @eval function Base.isapprox(o1::$Geo, o2::$Geo; kw...)
        isapprox(pga(o1), pga(o2); kw...)
    end
end

meet(o1::Plane, o2::Plane) = Line( pga(o1) ∧ pga(o2))
meet(o1::Line, o2::Plane)  = Point(pga(o1) ∧ pga(o2))
meet(o1::Plane, o2::Line)  = Point(pga(o1) ∧ pga(o2))

# TODO join
join(o::Union{Point, Line, Plane}) = o
join(o1::Point, o2::Point) = Line(pga(o1) ∨ pga(o2))
join(o1::Point, o2::Line) = Plane(pga(o1) ∨ pga(o2))
join(o1::Line, o2::Point) = Plane(pga(o1) ∨ pga(o2))
join(o1::Point, o2::Point, o3::Point) = Plane(pga(o1) ∨ pga(o2) ∨ pga(o3))

orthogonal_through(orth::Plane, through::Line) = Plane(inner(orth, through))
orthogonal_through(orth::Plane, through::Point) = Line(inner(orth, through))
orthogonal_through(orth::Line, through::Point) = Plane(inner(orth, through))

project(onto::Union{Point, Line, Plane}, obj::Point) = Point(PGA.project(pga(onto), pga(obj)))
project(onto::Plane, obj::Line) = Line(PGA.project(pga(onto), pga(obj)))

################################################################################
##### metric
################################################################################
function cos_angle(o1::Plane, o2::Plane)
    inner(pga(o1), pga(o2)).e
end
function cos_angle(o1::Line, o2::Plane)
    inner(pga(o1), pga(o2)).e
end
function cos_angle(o1::Plane, o2::Line)
    inner(pga(o1), pga(o2)).e
end

function distance(o1::Point, o2::Point)
    norm(pga(o1) ∨ pga(o2))
end

end #module