using VectorSphericalHarmonics
using Test
using Aqua
using LinearAlgebra
using LegendrePolynomials
using OffsetArrays
using SphericalHarmonics
using SphericalHarmonics: NorthPole, SouthPole, getY, eltypeY
using SphericalHarmonicModes
using WignerD
using HCubature
using VectorSphericalHarmonics: basisconversionmatrix, cache, cache!
using Rotations
using StaticArrays

@testset "project quality" begin
    if VERSION >= v"1.6"
        Aqua.test_all(VectorSphericalHarmonics, ambiguities = (recursive = false,))
    else
        Aqua.test_all(VectorSphericalHarmonics, ambiguities = false)
    end
end

isapproxdefault(a, b) = isapprox(a, b, atol = 1e-14, rtol = sqrt(eps(Float64)))
isapproxdefault(x::Tuple{Any, Any}) = isapproxdefault(x...)

function genspharm2(l, m, n, θ, ϕ)
    @assert abs(m) <= l "m must satisfy -l <= m <= l, received l = $l, m = $m"
    @assert abs(n) <= l "n must satisfy -l <= n <= l, received l = $l, n = $n"
    √((2l+1)/4π) * WignerD.wignerdjmn(l, m, n, θ) * cis(m * ϕ)
end

const S = cache(10);
const Ylm = getY(S);

@testset "Basis" begin
    for θ in LinRange(0, pi, 10), ϕ in LinRange(0, 2pi, 10)
        for B2 in [HelicityCovariant(), Polar(), SphericalCovariant(), Cartesian()]
            for B1 in [HelicityCovariant(), SphericalCovariant()]
                M = basisconversionmatrix(B1, B2, θ, ϕ)
                Minv = basisconversionmatrix(B2, B1, θ, ϕ)
                # transformations must be unitary
                @test M' * M ≈ I
                @test Minv ≈ M'
                @test abs(det(M)) ≈ 1
            end
        end
    end
    @testset "multiple transformations" begin
        for θ in LinRange(0, pi, 10), ϕ in LinRange(0, 2pi, 10)
            for B in [HelicityCovariant(), Polar(), SphericalCovariant(), Cartesian()]
                M1 = basisconversionmatrix(HelicityCovariant(), B, θ, ϕ)
                M2 = basisconversionmatrix(B, SphericalCovariant(), θ, ϕ)
                M3 = basisconversionmatrix(HelicityCovariant(), SphericalCovariant(), θ, ϕ)
                @test begin
                    res = M2 * M1 ≈ M3
                    if !res
                        println(HelicityCovariant(),"->",B,"->",SphericalCovariant(),", θ = $θ, ϕ = $ϕ")
                    end
                    res
                end

                M1 = basisconversionmatrix(SphericalCovariant(), B, θ, ϕ)
                M2 = basisconversionmatrix(B, HelicityCovariant(), θ, ϕ)
                M3 = basisconversionmatrix(SphericalCovariant(), HelicityCovariant(), θ, ϕ)
                @test begin
                    res = M2 * M1 ≈ M3
                    if !res
                        println(SphericalCovariant(),"->",B,"->",HelicityCovariant(),", θ = $θ, ϕ = $ϕ")
                    end
                    res
                end
            end
        end
    end
    @testset "basis at poles" begin
        @testset "Cartesian and polar" begin
            M = basisconversionmatrix(Cartesian(), Polar(), 0, 0)
            @test M * [1,0,0] ≈ [0,1,0] # x̂ == θ̂ at the north pole
            @test M * [0,1,0] ≈ [0,0,1] # ŷ == ϕ̂ at the north pole
            @test M * [0,0,1] ≈ [1,0,0] # ẑ == r̂ at the north pole

            M = basisconversionmatrix(Polar(), Cartesian(), pi, 0)
            @test M * [1,0,0] ≈ [0,0,-1] # r̂ == -ẑ at the south pole
            @test M * [0,1,0] ≈ [-1,0,0] # θ̂ == -x̂ at the south pole
            @test M * [0,0,1] ≈ [0,1,0]  # ϕ̂ == ŷ at the south pole
        end
    end
end

@testset "cache" begin
    θ, ϕ = pi/3, pi/4
    S2 = @inferred cache(Float64, 1)
    cache!(S2, θ, ϕ)
    @test getY(S2) == getY(cache(Float64, θ, ϕ, 1))

    S2 = @inferred cache(1)
    cache!(S2, θ, ϕ)
    @test getY(S2) == getY(cache(Float64, θ, ϕ, 1))

    @test vshbasis(PB(), Polar(), 1, 1, θ, ϕ, S2) == vshbasis(PB(), Polar(), 1, 1, θ, ϕ)
    @test genspharm(1, 1, θ, ϕ, S2) == genspharm(1, 1, θ, ϕ)
end

@testset "VSH orientations" begin
    θ1, ϕ1 = pi/3, pi/3;
    θNP, ϕNP = 0, 0;
    @testset "PB" begin
        @testset "HelicityCovariant" begin
            cache!(S, θ1, ϕ1, 0)
            Y = @inferred vshbasis(PB(), HelicityCovariant(), 0, 0, θ1, ϕ1, S)
            # for l=0 there is only one non-zero vector Ylm0, which is radial
            for n in axes(Y,2), basisind in axes(Y,1)
                if !(n == 0 && basisind == 0 #= e₀′ (radial) component =#)
                    @test isapproxdefault(0, Y[basisind, n])
                end
            end
        end
        @testset "Polar" begin
            cache!(S, θ1, ϕ1, 0)
            Y = @inferred vshbasis(PB(), Polar(), 0, 0, θ1, ϕ1, S)
            # for l=0 there is only one non-zero vector Ylm0, which is radial
            for n in axes(Y,2), basisind in axes(Y,1)
                if !(n == 0 && basisind == 1 #= radial component =#)
                    @test isapproxdefault(0, Y[basisind, n])
                end
            end
        end
        @testset "SphericalCovariant" begin
            # north pole
            cache!(S, θNP, ϕNP, 0)
            Y = @inferred vshbasis(PB(), SphericalCovariant(), 0, 0, θNP, ϕNP, S)
            # for l=0 there is only one non-zero vector Ylm0, which is radial
            # at north pole r̂ == ẑ. so there is only one non-zero component of Ylm0
            for n in axes(Y,2), basisind in axes(Y,1)
                if !(n == 0 && basisind == 0 #= e₀ (z) component =#)
                    @test isapproxdefault(0, Y[basisind, n])
                end
            end
        end
        @testset "Cartesian" begin
            # north pole
            cache!(S, θNP, ϕNP, 0)
            Y = @inferred vshbasis(PB(), Cartesian(), 0, 0, θNP, ϕNP, S)
            # for l=0 there is only one non-zero vector Ylm0, which is radial
            # at north pole r̂ == ẑ. so there is only one non-zero component of Ylm0
            for n in axes(Y,2), basisind in axes(Y,1)
                if !(n == 0 && basisind == 3 #= z component =#)
                    @test isapproxdefault(0, Y[basisind, n])
                end
            end
        end
    end
    @testset "Hansen" begin
        @testset "HelicityCovariant" begin
            cache!(S, θ1, ϕ1, 0)
            Y = @inferred vshbasis(Hansen(), HelicityCovariant(), 0, 0, θ1, ϕ1, S)
            # for l=0 there is only one non-zero vector Ylm⁽⁻¹⁾, which is radial
            for n in axes(Y,2), basisind in axes(Y,1)
                if !(n == -1 && basisind == 0 #= e₀′ (radial) component =#)
                    @test isapproxdefault(0, Y[basisind, n])
                end
            end
        end
        @testset "Polar" begin
            cache!(S, θ1, ϕ1, 0)
            Y = @inferred vshbasis(Hansen(), Polar(), 0, 0, θ1, ϕ1, S)
            # for l=0 there is only one non-zero vector Ylm⁽⁻¹⁾, which is radial
            for n in axes(Y,2), basisind in axes(Y,1)
                if !(n == -1 && basisind == 1 #= radial component =#)
                    @test isapproxdefault(0, Y[basisind, n])
                end
            end
        end
        @testset "SphericalCovariant" begin
            # north pole
            cache!(S, θNP, ϕNP, 0)
            Y = @inferred vshbasis(Hansen(), SphericalCovariant(), 0, 0, θNP, ϕNP, S)
            # for l=0 there is only one non-zero vector Ylm⁽⁻¹⁾, which is radial
            # at north pole r̂ == ẑ. so there is only one non-zero component of Ylm0
            for n in axes(Y,2), basisind in axes(Y,1)
                if !(n == -1 && basisind == 0 #= e₀ (z) component =#)
                    @test isapproxdefault(0, Y[basisind, n])
                end
            end
        end
        @testset "Cartesian" begin
            # north pole
            cache!(S, θNP, ϕNP, 0)
            Y = @inferred vshbasis(Hansen(), Cartesian(), 0, 0, θNP, ϕNP, S)
            # for l=0 there is only one non-zero vector Ylm⁽⁻¹⁾⁾, which is radial
            # at north pole r̂ == ẑ. so there is only one non-zero component of Ylm0
            for n in axes(Y,2), basisind in axes(Y,1)
                if !(n == -1 && basisind == 3 #= z component =#)
                    @test isapproxdefault(0, Y[basisind, n])
                end
            end
        end
    end
    @testset "Irreducible" begin
        cache!(S, θ1, ϕ1, 0)
        Y = @inferred vshbasis(Irreducible(), Polar(), 0, 0, θ1, ϕ1, S)
        @test all(iszero, Y[:, -1:0])
        @test all(x -> isapproxdefault(x, 0), Y[2:3, 1])
        @test Y[1, 1] ≈ -1/√(4pi)

        for m = -1:1
            Y = @inferred vshbasis(Irreducible(), SphericalCovariant(), 1, m, θ1, ϕ1, S)
            @test Y[m, 0] ≈ 1/√(4pi)
            @test all(iszero, Y[-1:m-1, 0])
            @test all(iszero, Y[m+1:1, 0])
        end

        cache!(S, θNP, ϕNP, 0)
        Y = @inferred vshbasis(Irreducible(), Cartesian(), 0, 0, θNP, ϕNP, S)
        @test all(iszero, Y[:, -1:0])
        @test all(iszero, Y[1:2, 1])
        @test Y[3, 1] ≈ -1/√(4pi)
    end
end

@testset "norm in different bases" begin
    lmax = 4
    for θ in LinRange(0, pi, 10), ϕ in LinRange(0, 2pi, 10)
        cache!(S, θ, ϕ)
        for l = 0:lmax, m = -l:l
            Y = vshbasis(Irreducible(), SphericalCovariant(), l, m, θ, ϕ, S)
            Yp = parent(Y)
            for B in [HelicityCovariant(), Polar(), Cartesian()]
                Y2 = vshbasis(Irreducible(), B, l, m, θ, ϕ, S)
                Y2p = parent(Y2)
                for l in 1:3
                    @test norm(Y2p[:, l]) ≈ norm(Yp[:, l])
                end
            end
        end
    end
end

@testset "VSH components" begin
    lmax = 5
    @testset "HelicityCovariant" begin
        @testset "PB" begin
            for θ in LinRange(0, pi, 10), ϕ in LinRange(0, 2pi, 10)
                cache!(S, θ, ϕ)
                for l in 0:lmax, m in -l:l
                    Y = vshbasis(PB(), HelicityCovariant(), l, m, θ, ϕ, S)
                    @test isapproxdefault(Y[0,0], Ylm[(l,m)])
                    for n in -1:1
                        if n in -min(1, l):min(1, l)
                            @test isapproxdefault(Y[n,n], genspharm2(l, m, n, θ, ϕ))
                        else
                            @test isapproxdefault(Y[n,n], 0)
                        end
                        for basisind in -1:1
                            basisind == n && continue
                            @test isapproxdefault(Y[basisind,n], 0)
                        end
                    end
                end
            end
        end
        @testset "Hansen" begin
            # l = 0 has only one component
            for θ in LinRange(0, pi, 10), ϕ in LinRange(0, 2pi, 10)
                cache!(S, θ, ϕ, 0)
                Y = vshbasis(Hansen(), HelicityCovariant(), 0, 0, θ, ϕ, S)
                @test isapproxdefault(Y[ 0, -1],  √(1/4pi) * WignerD.wignerDjmn(0, 0, 0, 0, θ, ϕ))
                for n in -1:1, basisind in -1:1
                    basisind == 0 && n == -1 && continue
                    @test isapproxdefault(Y[basisind, n], 0)
                end
            end
            for θ in LinRange(0, pi, 10), ϕ in LinRange(0, 2pi, 10)
                cache!(S, θ, ϕ)
                for l in 1:lmax, m in -l:l
                    norm = √((2l+1)/8pi)
                    sqrt2norm = √((2l+1)/4pi)
                    Y = vshbasis(Hansen(), HelicityCovariant(), l, m, θ, ϕ, S)
                    @test isapproxdefault(Y[-1, +1],  norm * WignerD.wignerDjmn(l,  1, -m, 0, θ, ϕ))
                    @test isapproxdefault(Y[ 0, +1],  0)
                    @test isapproxdefault(Y[+1, +1],  norm * WignerD.wignerDjmn(l, -1, -m, 0, θ, ϕ))
                    @test isapproxdefault(Y[-1,  0],  norm * WignerD.wignerDjmn(l,  1, -m, 0, θ, ϕ))
                    @test isapproxdefault(Y[ 0,  0],  0)
                    @test isapproxdefault(Y[+1,  0], -norm * WignerD.wignerDjmn(l, -1, -m, 0, θ, ϕ))
                    @test isapproxdefault(Y[-1, -1],  0)
                    @test isapproxdefault(Y[ 0, -1],  sqrt2norm * WignerD.wignerDjmn(l, 0, -m, 0, θ, ϕ))
                    @test isapproxdefault(Y[+1, -1],  0)
                end
            end
        end
        @testset "Irreducible" begin
            # j = 0 has only one component
            for θ in LinRange(0, pi, 10), ϕ in LinRange(0, 2pi, 10)
                cache!(S, θ, ϕ)
                Y = vshbasis(Irreducible(), HelicityCovariant(), 0, 0, θ, ϕ, S)
                @test isapproxdefault(Y[0, +1], -√(1/4pi) * WignerD.wignerDjmn(0, 0, 0, 0, θ, ϕ))
                for n in -1:1, basisind in -1:1
                    basisind == 0 && n == 1 && continue
                    @test isapproxdefault(Y[basisind, n], 0)
                end
            end
            for θ in LinRange(0, pi, 10), ϕ in LinRange(0, 2pi, 10)
                cache!(S, θ, ϕ)
                for j in 1:lmax, m in -j:j
                    Y = vshbasis(Irreducible(), HelicityCovariant(), j, m, θ, ϕ, S)
                    @test isapproxdefault(Y[-1, j+1], √(j/8pi) * WignerD.wignerDjmn(j,  1, -m, 0, θ, ϕ))
                    @test isapproxdefault(Y[ 0, j+1], -√((j+1)/4pi) * WignerD.wignerDjmn(j, 0, -m, 0, θ, ϕ))
                    @test isapproxdefault(Y[+1, j+1], √(j/8pi) * WignerD.wignerDjmn(j, -1, -m, 0, θ, ϕ))
                    @test isapproxdefault(Y[-1,  j], √((2j+1)/8pi) * WignerD.wignerDjmn(j,  1, -m, 0, θ, ϕ))
                    @test isapproxdefault(Y[ 0,  j], 0)
                    @test isapproxdefault(Y[+1,  j], -√((2j+1)/8pi) * WignerD.wignerDjmn(j, -1, -m, 0, θ, ϕ))
                    @test isapproxdefault(Y[-1, j-1],  √((j+1)/8pi) * WignerD.wignerDjmn(j,  1, -m, 0, θ, ϕ))
                    @test isapproxdefault(Y[ 0, j-1],  √(j/4pi) * WignerD.wignerDjmn(j, 0, -m, 0, θ, ϕ))
                    @test isapproxdefault(Y[+1, j-1],  √((j+1)/8pi) * WignerD.wignerDjmn(j, -1, -m, 0, θ, ϕ))
                end
            end
        end
    end
    @testset "Polar" begin
        @testset "Hansen" begin
            # l = 0 has only one component
            for θ in LinRange(0, pi, 10), ϕ in LinRange(0, 2pi, 10)
                cache!(S, θ, ϕ, 0)
                Y = vshbasis(Hansen(), Polar(), 0, 0, θ, ϕ, S)
                @test begin
                    res = isapproxdefault(Y[ 1, -1], Ylm[(0,0)])
                    if !res
                        @show θ, ϕ
                    end
                    res
                end
                for n in axes(Y,2), basisind in axes(Y,1)
                    basisind == 1 && n == -1 && continue
                    @test isapproxdefault(Y[basisind, n], 0)
                end
            end
            for θ in LinRange(0, pi, 10), ϕ in LinRange(0, 2pi, 10)
                cache!(S, θ, ϕ)
                for l in 1:lmax, m in -l:l
                    norml⁻ = 1/2*√(((l-m)*(l+m+1))/(l*(l+1)))
                    norml⁺ = 1/2*√(((l+m)*(l-m+1))/(l*(l+1)))
                    Y = vshbasis(Hansen(), Polar(), l, m, θ, ϕ, S)
                    @test isapproxdefault(Y[1, +1],  0)
                    @test isapproxdefault(Y[2, +1],  (m < l ? norml⁻*cis(-ϕ)*Ylm[(l, m+1)] : 0) -
                                          (m > -l ? norml⁺*cis(ϕ)*Ylm[(l, m-1)] : 0)
                        )
                    if sin(θ) != 0
                        @test isapproxdefault(Y[3, +1],  im*m/√(l*(l+1))*1/sin(θ)*Ylm[(l,m)])
                    end
                    @test isapproxdefault(Y[1,  0],  0)
                    if sin(θ) != 0
                        @test isapproxdefault(Y[2,  0],  -m/√(l*(l+1))*1/sin(θ)*Ylm[(l,m)])
                    end
                    @test isapproxdefault(Y[3, 0],  (m < l ? -im*norml⁻*cis(-ϕ)*Ylm[(l, m+1)] : 0) +
                                          (m > -l ? im*norml⁺*cis(ϕ)*Ylm[(l, m-1)] : 0)
                        )
                    @test isapproxdefault(Y[1, -1],  Ylm[(l,m)])
                    @test isapproxdefault(Y[2, -1],  0)
                    @test isapproxdefault(Y[3, -1],  0)

                    if m != 0 && θ == 0
                        # n.Y⁽¹⁾ and n.Y⁽⁰⁾ are zero, and n.Y⁽⁻¹⁾ == YLM which is zero at poles for M != 0
                        @test all(x -> isapproxdefault(x, 0), Y[1, :])
                    end
                end
            end
        end
        @testset "Irreducible" begin
            @testset "Pole" begin
                θ = 0; ϕ = 0
                cache!(S, θ, ϕ)
                for l in 1:lmax, m in -l:l
                    Y = vshbasis(Irreducible(), Polar(), l, m, θ, ϕ, S)
                    # n.Yᴶ⁺¹ and n.Yᴶ⁻¹ ∝ YLM which is zero at poles for M != 0. n.Yᴶ is zero
                    m == 0 && continue
                    @test all(x -> isapproxdefault(x, 0), Y[1, :])
                end
            end
            @testset "J = 0, M = 0" begin
                for θ in LinRange(0, pi, 10), ϕ in LinRange(0, 2pi, 10)
                    cache!(S, θ, ϕ)
                    Y = vshbasis(Irreducible(), Polar(), 0, 0, θ, ϕ, S)
                    @test all(iszero, Y[:,-1:0])
                    @test Y[1, 1] ≈ -1/√(4pi)
                    @test begin
                        res = all(x -> isapproxdefault(0, x), Y[2:3, 1])
                        if !res
                            @show θ, ϕ
                        end
                        res
                    end
                end
            end
        end
    end
    @testset "SphericalCovariant" begin
        @testset "Irreducible" begin
            # For l = 0 only Y₀₀¹ is non-zero
            for θ in LinRange(0, pi, 10), ϕ in LinRange(0, 2pi, 10)
                cache!(S, θ, ϕ, 0)
                Y = vshbasis(Irreducible(), SphericalCovariant(), 0, 0, θ, ϕ, S)
                @test all(iszero, Y[:, -1:0])
                @test isapproxdefault(Y[1, 1], 1/√3 * Ylm[(1,-1)])
                @test isapproxdefault(Y[0, 1], -1/√3 * Ylm[(1, 0)])
                @test isapproxdefault(Y[-1, 1], 1/√3 * Ylm[(1, 1)])
            end
            for θ in LinRange(0, pi, 10), ϕ in LinRange(0, 2pi, 10)
                cache!(S, θ, ϕ)
                for j in 1:lmax, m in -j:j
                    Y = vshbasis(Irreducible(), SphericalCovariant(), j, m, θ, ϕ, S)
                    @test isapproxdefault(Y[ 1, j+1],  √((j-m+1)*(j-m+2)/(2*(j+1)*(2j+3))) * Ylm[(j+1,m-1)])
                    @test isapproxdefault(Y[ 0, j+1], -√((j-m+1)*(j+m+1)/((j+1)*(2j+3)))   * Ylm[(j+1,m)])
                    @test isapproxdefault(Y[-1, j+1],  √((j+m+1)*(j+m+2)/(2*(j+1)*(2j+3))) * Ylm[(j+1,m+1)])
                    @test isapproxdefault(Y[ 1, j], m > -j ? -√((j+m)*(j-m+1)/(2j*(j+1))) * Ylm[(j,m-1)] : 0)
                    @test isapproxdefault(Y[ 0, j],  m/√(j*(j+1)) * Ylm[(j,m)])
                    @test isapproxdefault(Y[-1, j], m < j ? √((j-m)*(j+m+1)/(2j*(j+1))) * Ylm[(j,m+1)] : 0)
                    @test isapproxdefault(Y[ 1, j-1], abs(m - 1) <= j-1 ? √((j+m)*(j+m-1)/(2j*(2j-1))) * Ylm[(j-1,m-1)] : 0)
                    @test isapproxdefault(Y[ 0, j-1], abs(m) <= j-1 ? √((j-m)*(j+m)/(j*(2j-1))) * Ylm[(j-1,m)] : 0)
                    @test isapproxdefault(Y[-1, j-1], abs(m+1) <= j-1 ? √((j-m)*(j-m-1)/(2j*(2j-1))) * Ylm[(j-1,m+1)] : 0)
                end
            end
            @testset "Pole" begin
                @testset "NorthPole" begin
                    cache!(S, NorthPole(), 0, lmax)
                    for j in 1:lmax, m in -j:j
                        # test that θ = 0 and θ = NorthPole() lead to identical results
                        Y0 = vshbasis(Irreducible(), SphericalCovariant(), j, m, 0, 0)
                        Y = vshbasis(Irreducible(), SphericalCovariant(), j, m, NorthPole(), 0, S)
                        @test isapproxdefault(Y0, Y)
                        if abs(m) > 1
                            @test all(iszero, Y)
                        elseif m == 1 || m == -1
                            @test Y[m,j+1] ≈ √(j/8pi)
                            @test Y[m,j] ≈ -m*√((2j+1)/8pi)
                            @test Y[m,j-1] ≈ √((j+1)/8pi)
                        elseif m == 0
                            @test Y[0,j+1] ≈ -√((j+1)/4pi)
                            @test isapproxdefault(Y[0,j], 0)
                            @test Y[0,j-1] ≈ √(j/4pi)
                        end

                        Y = vshbasis(Hansen(), SphericalCovariant(), j, m, NorthPole(), 0, S)
                        if abs(m) > 1
                            @test all(iszero, Y)
                        elseif m == 1 || m == -1
                            @test Y[m, 1] ≈ √((2j+1)/8pi)
                            @test Y[m, 0] ≈ -m*√((2j+1)/8pi)
                            @test isapproxdefault(Y[m,-1], 0)
                        elseif m == 0
                            @test isapproxdefault(Y[0, 1], 0)
                            @test isapproxdefault(Y[0, 0], 0)
                            @test Y[0,-1] ≈ √((2j+1)/4pi)
                        end
                    end
                end
                @testset "SouthPole" begin
                    cache!(S, SouthPole(), 0, lmax)
                    for j in 1:lmax, m in -j:j
                        # test that θ = pi and θ = SouthPole() lead to identical results
                        Y0 = vshbasis(Irreducible(), SphericalCovariant(), j, m, pi, 0)
                        Y = vshbasis(Irreducible(), SphericalCovariant(), j, m, SouthPole(), 0, S)
                        @test isapproxdefault(Y0, Y)
                        if abs(m) > 1
                            @test all(iszero, Y)
                        elseif m == 1 || m == -1
                            @test Y[m,j+1] ≈ (-1)^(j-1)*√(j/8pi)
                            @test Y[m,j] ≈ -(-1)^(j)*m*√((2j+1)/8pi)
                            @test Y[m,j-1] ≈ (-1)^(j-1)*√((j+1)/8pi)
                        elseif m == 0
                            @test Y[0,j+1] ≈ -(-1)^(j-1)*√((j+1)/4pi)
                            @test isapproxdefault(Y[0,j], 0)
                            @test Y[0,j-1] ≈ (-1)^(j-1)*√(j/4pi)
                        end

                        Y = vshbasis(Hansen(), SphericalCovariant(), j, m, SouthPole(), 0, S)
                        if abs(m) > 1
                            @test all(iszero, Y)
                        elseif m == 1 || m == -1
                            @test Y[m, 1] ≈ (-1)^(j-1)*√((2j+1)/8pi)
                            @test Y[m, 0] ≈ -(-1)^(j)*m*√((2j+1)/8pi)
                            @test isapproxdefault(Y[m,-1], 0)
                        elseif m == 0
                            @test isapproxdefault(Y[0, 1], 0)
                            @test isapproxdefault(Y[0, 0], 0)
                            @test Y[0,-1] ≈ (-1)^(j-1)*√((2j+1)/4pi)
                        end
                    end
                end
            end
        end
        @testset "Hansen" begin
            for θ in LinRange(0, pi, 10), ϕ in LinRange(0, 2pi, 10)
                cache!(S, θ, ϕ, 0)
                Y = vshbasis(Hansen(), SphericalCovariant(), 0, 0, θ, ϕ, S)
                @test all(iszero, Y[:, 0:1])
                @test isapproxdefault(Y[ 1, -1], -1/√3 * Ylm[(1,-1)])
                @test isapproxdefault(Y[ 0, -1], 1/√3 * Ylm[(1, 0)])
                @test isapproxdefault(Y[-1, -1], -1/√3 * Ylm[(1, 1)])
            end
            for θ in LinRange(0, pi, 10), ϕ in LinRange(0, 2pi, 10)
                cache!(S, θ, ϕ)
                for l in 1:lmax, m in -l:l
                    Y = vshbasis(Hansen(), SphericalCovariant(), l, m, θ, ϕ, S)
                    @test isapproxdefault(Y[ 1, 1],  √((l-m+1)*(l-m+2)*l/(2*(l+1)*(2l+1)*(2l+3))) * Ylm[(l+1,m-1)] +
                                                    (abs(m - 1) <= l-1 ? √((l+m)*(l+m-1)*(l+1)/(2l*(2l-1)*(2l+1))) * Ylm[(l-1,m-1)] : 0))
                    @test isapproxdefault(Y[ 0, 1], -√((l-m+1)*(l+m+1)*l/((l+1)*(2l+1)*(2l+3)))   * Ylm[(l+1,m)] +
                                                    (abs(m) <= l-1 ? √((l-m)*(l+m)*(l+1)/(l*(2l-1)*(2l+1))) * Ylm[(l-1,m)] : 0))
                    @test isapproxdefault(Y[-1, 1],  √((l+m+1)*(l+m+2)*l/(2*(l+1)*(2l+1)*(2l+3))) * Ylm[(l+1,m+1)] +
                                                    (abs(m+1) <= l-1 ? √((l-m)*(l-m-1)*(l+1)/(2l*(2l-1)*(2l+1))) * Ylm[(l-1,m+1)] : 0))
                    @test isapproxdefault(Y[ 1, 0], m > -l ? -√((l+m)*(l-m+1)/(2l*(l+1))) * Ylm[(l,m-1)] : 0)
                    @test isapproxdefault(Y[ 0, 0],  m/√(l*(l+1)) * Ylm[(l,m)])
                    @test isapproxdefault(Y[-1, 0], m < l ? √((l-m)*(l+m+1)/(2l*(l+1))) * Ylm[(l,m+1)] : 0)
                    @test isapproxdefault(Y[ 1,-1], (abs(m - 1) <= l-1 ? √((l+m)*(l+m-1)/(2*(2l-1)*(2l+1))) * Ylm[(l-1,m-1)] : 0) -
                                                    √((l-m+1)*(l-m+2)/(2*(2l+1)*(2l+3))) * Ylm[(l+1,m-1)])
                    @test isapproxdefault(Y[ 0,-1], (abs(m) <= l-1 ? √((l-m)*(l+m)/((2l-1)*(2l+1))) * Ylm[(l-1,m)] : 0) +
                                                    √((l-m+1)*(l+m+1)/((2l+1)*(2l+3))) * Ylm[(l+1,m)])
                    @test isapproxdefault(Y[-1,-1], (abs(m+1) <= l-1 ? √((l-m)*(l-m-1)/(2*(2l-1)*(2l+1))) * Ylm[(l-1,m+1)] : 0) -
                                                    √((l+m+1)*(l+m+2)/(2*(2l+1)*(2l+3))) * Ylm[(l+1,m+1)])
                end
            end
        end
    end
    @testset "matrix and vectors" begin
        for θ in LinRange(0, pi, 10), ϕ in LinRange(0, 2pi, 10)
            cache!(S, θ, ϕ)
            for YT in [PB(), Irreducible(), Hansen()],
                B in [Cartesian(), Polar(), SphericalCovariant(), HelicityCovariant()]

                for l in 0:lmax, m in -l:l
                    Y = vshbasis(YT, B, l, m, θ, ϕ, S)
                    Yp = parent(Y)
                    for n in axes(Y, 2)
                        Yn = vshbasis(YT, B, l, m, n, θ, ϕ, S)
                        @test begin
                            res = parent(Y[:, n]) ≈ parent(Yn)
                            if !res
                                @show YT, B, l, m, n
                            end
                            res
                        end
                    end
                end
            end
        end

        @testset "genspharm" begin
            for θ in LinRange(0, pi, 10), ϕ in LinRange(0, 2pi, 10)
                cache!(S, θ, ϕ)
                G2 = genspharm(ML(0:lmax), θ, ϕ, S)
                for l in 0:lmax, m in -l:l
                    Y = vshbasis(PB(), HelicityCovariant(), l, m, θ, ϕ, S)
                    G = genspharm(l, m, θ, ϕ, S)
                    @test isapproxdefault(G2[(l,m)], G)
                    nmax = min(1, l)
                    for n in -nmax:nmax
                        Ylmn = genspharm2(l, m, n, θ, ϕ)
                        @test isapproxdefault(Y[n,n], G[n])
                        @test isapproxdefault(G[n], Ylmn)
                    end
                end
            end
        end
    end
end

@testset "one vs multiple" begin
    lmax = 2
    modes = LM(0:lmax)
    modes_high = LM(2:lmax, 2:2)
    S2 = VectorSphericalHarmonics.cache(Float64, lmax)
    for YT in [Irreducible(), PB(), Hansen()], B in [Cartesian(), Polar(), SphericalCovariant(), HelicityCovariant()]
        @testset "$YT $B" begin
            for m in [modes, modes_high], (θ, ϕ) in ((pi/3, pi/3), (0,0))
                cache!(S2, θ, ϕ)
                Y_multiple = vshbasis(YT, B, m, θ, ϕ, S2)
                Y_multiple2 = vshbasis(YT, B, m, θ, ϕ)
                for (ind, (j, m)) in enumerate(m)
                    Y = vshbasis(YT, B, j, m, θ, ϕ, S2)
                    @test isapproxdefault(Y_multiple[ind], Y)
                    @test isapproxdefault(Y_multiple2[ind], Y)
                end
            end
        end
    end

    @testset "genspharm" begin
        for m in [modes, modes_high], (θ, ϕ) in ((pi/3, pi/3), (0,0))
            cache!(S2, θ, ϕ)
            Y_multiple = genspharm(m, θ, ϕ, S2)
            Y_multiple2 = genspharm(m, θ, ϕ)
            for (ind, (j, m)) in enumerate(m)
                Y = genspharm(j, m, θ, ϕ, S2)
                @test isapproxdefault(Y_multiple[ind], Y)
                @test isapproxdefault(Y_multiple2[ind], Y)
            end
        end
    end
end

@testset "Orthogonal" begin
    lmax = 10
    @testset "Locally orthogonal" begin
        @testset "PB" begin
            # The PB basis vectors are directed along the Helicity basis vectors, so they are orthogonal at each point
            for θ in LinRange(0, pi, 10), ϕ in LinRange(0, 2pi, 10)
                cache!(S, θ, ϕ);
                for l in 0:lmax, m in -l:l, B in [SphericalCovariant(), Polar(), HelicityCovariant(), Cartesian()]
                    Y = vshbasis(PB(), B, l, m, θ, ϕ, S)
                    Yp = parent(Y)
                    M = Yp'*Yp
                    for j in axes(M, 2), i in axes(M, 1)
                        i == j && continue
                        @test isapprox(M[i, j], 0, atol=1e-13, rtol=1e-8)
                    end
                end
            end
        end
        @testset "Hansen" begin
            # For Hansen basis, the local orthogonality uses transpose instead of adjoint.
            for θ in LinRange(0, pi, 10), ϕ in LinRange(0, 2pi, 10)
                cache!(S, θ, ϕ);
                for l in 0:lmax, m in -l:l
                    for B in [Polar(), Cartesian()]
                        Y = vshbasis(Hansen(), B, l, m, θ, ϕ, S)
                        Yp = parent(Y)
                        M = transpose(Yp)*Yp
                        for j in axes(M, 2), i in axes(M, 1)
                            i == j && continue
                            @test begin
                                res = isapprox(M[i, j], 0, atol=1e-13, rtol=1e-8)
                                if !res
                                    @show B, l, m, θ, ϕ
                                end
                                res
                            end
                        end
                    end
                    # For the complex bases, we need to use eμ⋅eν = (-1)^μ δμ,-ν
                    for B in [SphericalCovariant(), HelicityCovariant()]
                        Y = vshbasis(Hansen(), B, l, m, θ, ϕ, S)
                        for j in axes(Y, 2), i in axes(Y, 2)
                            i == j && continue
                            M = sum(Y[b,j] * Y[-b,i] * (-1)^b for b = -1:1)
                            @test begin
                                res = isapprox(M, 0, atol=1e-13, rtol=1e-8)
                                if !res
                                    @show B, l, m, θ, ϕ
                                end
                                res
                            end
                        end
                    end
                end
            end
        end
    end
    @testset "Orthonormal" begin
        lmax = 3
        function f(θ, ϕ, YT, B, l, m)
            cache!(S, θ, ϕ)
            Y = vshbasis(YT, B, l, m, θ, ϕ, S)
            Yp = parent(Y)
            (Yp'*Yp) .* sin(θ)
        end
        @testset "PB" begin
            for l in 0:lmax, m in -l:l, B in [SphericalCovariant(), Polar(), HelicityCovariant(), Cartesian()]
                M = hcubature(x -> f(x..., PB(), B, l, m), [0, 0], [pi, 2pi])[1]
                if l == 0
                    for I in CartesianIndices(M)
                        i, j = Tuple(I)
                        i == 2 && j == 2 && continue
                        @test isapproxdefault(M[i,j], 0)
                    end
                    @test M[2,2] ≈ 1
                else
                    @test M ≈ I
                end
            end
        end
        @testset "Irreducible" begin
            for l in 0:lmax, m in -l:l, B in [SphericalCovariant(), Polar(), HelicityCovariant(), Cartesian()]
                M = hcubature(x -> f(x..., Irreducible(), B, l, m), [0, 0], [pi, 2pi])[1]
                if l == 0
                    for I in CartesianIndices(M)
                        i, j = Tuple(I)
                        i == 3 && j == 3 && continue
                        @test isapproxdefault(M[i,j], 0)
                    end
                    @test M[3,3] ≈ 1
                else
                    @test M ≈ I
                end
            end
        end
        @testset "Hansen" begin
            for l in 0:lmax, m in -l:l, B in [SphericalCovariant(), Polar(), HelicityCovariant(), Cartesian()]
                M = hcubature(x -> f(x..., Hansen(), B, l, m), [0, 0], [pi, 2pi])[1]
                if l == 0
                    for I in CartesianIndices(M)
                        i, j = Tuple(I)
                        i == 1 && j == 1 && continue
                        @test isapproxdefault(M[i,j], 0)
                    end
                    @test M[1,1] ≈ 1
                else
                    @test M ≈ I
                end
            end
        end
    end
end

@testset "Addition theorem" begin
    lmax = 10
    @testset "same point" begin
        @testset "Irreducible" begin
            θ, ϕ = pi/3, pi/3
            for B in [Polar(), SphericalCovariant(), Cartesian(), HelicityCovariant()]
                Y1_all = vshbasis(Irreducible(), B, ML(0:lmax), θ, ϕ);
                for j in 1:lmax, l1 in j-1:j+1, l2 in j-1:j+1
                    S = zero(Float64)
                    for m in -j:j
                        Yjm = Y1_all[(j,m)]
                        Yjml1 = parent(Yjm[:,l1])
                        Yjml2 = parent(Yjm[:,l2])
                        S += dot(Yjml1, Yjml2)
                    end
                    if l1 == l2
                        @test isapproxdefault(S, (2j+1)/4pi)
                    else
                        @test isapproxdefault(S, 0)
                    end
                end
            end
        end
        @testset "Hansen" begin
            θ, ϕ = pi/3, pi/3
            for B in [Polar(), SphericalCovariant(), Cartesian(), HelicityCovariant()]
                Y1_all = vshbasis(Hansen(), B, ML(0:lmax), θ, ϕ);
                for j in 1:lmax, λ1 in -1:+1, λ2 in -1:+1
                    S = zero(Float64)
                    for m in -j:j
                        Yjm = Y1_all[(j,m)]
                        Yjmλ1 = parent(Yjm[:,λ1])
                        Yjmλ2 = parent(Yjm[:,λ2])
                        S += dot(Yjmλ1, Yjmλ2)
                    end
                    if λ1 == λ2
                        @test isapproxdefault(S, (2j+1)/4pi)
                    else
                        @test isapproxdefault(S, 0)
                    end
                end
            end
        end
    end
    @testset "different points" begin
        @testset "Irreducible" begin
            θ1, ϕ1 = pi/3, pi/3
            θ2, ϕ2 = pi/4, pi/4
            cosω12 = cos(θ1)cos(θ2) + sin(θ1)sin(θ2)cos(ϕ1 - ϕ2)
            P = LegendrePolynomials.collectPl(cosω12, lmax = lmax + 1);
            allmodes = ML(0:lmax)
            # We only test in bases that are constant, so they are identical at both the points
            for B in [SphericalCovariant(), Cartesian()]
                Y1_all = vshbasis(Irreducible(), B, allmodes, θ1, ϕ1);
                Y2_all = vshbasis(Irreducible(), B, allmodes, θ2, ϕ2);
                for j in 1:lmax, l1 in j-1:j+1, l2 in j-1:j+1
                    x = zero(Float64)
                    for m in -j:j
                        Y1jm = Y1_all[(j,m)]
                        Y2jm = Y2_all[(j,m)]
                        Y1jml1 = parent(Y1jm[:,l1])
                        Y2jml2 = parent(Y2jm[:,l2])
                        x += dot(Y1jml1, Y2jml2)
                    end
                    @test isapproxdefault(imag(x), 0)
                    if l1 == l2
                        @test begin
                            res = isapproxdefault(x, (2j+1)/4pi * P[l1])
                            if !res
                                @show j, l1, B, x, (2j+1)/4pi * P[l1]
                            end
                            res
                        end
                    else
                        @test begin
                            res = isapproxdefault(x, 0)
                            if !res
                                @show j, l1, l2, B, x
                            end
                            res
                        end
                    end
                end
            end
        end
    end
end

@testset "Complex conjugation" begin
    @testset "Irreducible" begin
        # For l = 0 Y₀₀¹ is purely real
        for θ in LinRange(0, pi, 10), ϕ in LinRange(0, 2pi, 10)
            cache!(S, θ, ϕ, 0)
            Y = vshbasis(Irreducible(), Polar(), 0, 0, θ, ϕ, S)
            @test all(x -> isapproxdefault(imag(x), 0), Y)
        end
        for θ in LinRange(0, pi, 10), ϕ in LinRange(0, 2pi, 10)
            cache!(S, θ, ϕ)
            for j in 1:5, m in -j:j
                Ym = vshbasis(Irreducible(), Polar(), j, m, θ, ϕ, S)
                Ymconj = conj.(Ym)
                Y₋m = vshbasis(Irreducible(), Polar(), j, -m, θ, ϕ, S)
                for l in axes(Ym, 2)
                    Ymnconj = Ymconj[:, l]
                    Y₋mn = Y₋m[:, l]
                    phase = (-1)^(j + m + l + 1)
                    @test all(isapproxdefault, zip(Ymnconj, phase .* Y₋mn))
                end
            end
        end
    end
    @testset "Hansen" begin
        # For l = 0 Y₀₀⁽⁻¹⁾ is purely real
        for θ in LinRange(0, pi, 10), ϕ in LinRange(0, 2pi, 10)
            cache!(S, θ, ϕ, 0)
            Y = vshbasis(Hansen(), Polar(), 0, 0, θ, ϕ, S)
            @test all(x -> isapproxdefault(imag(x), 0), Y)
        end
        for θ in LinRange(0, pi, 10), ϕ in LinRange(0, 2pi, 10)
            cache!(S, θ, ϕ)
            for l in 1:5, m in -l:l
                Ym = vshbasis(Hansen(), Polar(), l, m, θ, ϕ, S)
                Ymconj = conj.(Ym)
                Y₋m = vshbasis(Hansen(), Polar(), l, -m, θ, ϕ, S)
                for λ in axes(Ym, 2)
                    Ymnconj = Ymconj[:, λ]
                    Y₋mn = Y₋m[:, λ]
                    phase = (-1)^(m + λ + 1)
                    @test all(isapproxdefault.(Ymnconj, phase .* Y₋mn))
                end
            end
        end
    end
    @testset "PB" begin
        for θ in LinRange(0, pi, 10), ϕ in LinRange(0, 2pi, 10)
            cache!(S, θ, ϕ)
            for j in 1:5, m in -j:j
                Ym = vshbasis(PB(), Polar(), j, m, θ, ϕ, S)
                Ymconj = conj.(Ym)
                Y₋m = vshbasis(PB(), Polar(), j, -m, θ, ϕ, S)
                for n in axes(Ym, 2)
                    Y₋m₋n = Y₋m[:, -n]
                    Ymn = Ym[:, n]
                    phase = (-1)^m
                    @test begin
                        res = all(isapproxdefault, zip(conj.(Y₋m₋n), phase .* Ymn))
                        if !res
                            @show j, m, n, θ, ϕ
                        end
                        res
                    end
                end
            end
        end
        for θ in LinRange(0, pi, 10), ϕ in LinRange(0, 2pi, 10)
            cache!(S, θ, ϕ)
            for j in 1:5, m in -j:j
                Ym = vshbasis(PB(), HelicityCovariant(), j, m, θ, ϕ, S)
                Ymconj = conj.(Ym)
                Y₋m = vshbasis(PB(), HelicityCovariant(), j, -m, θ, ϕ, S)
                for n in axes(Ym, 2)
                    Y₋m₋n = Y₋m[-n, -n]
                    Ymn = Ym[n, n]
                    phase = (-1)^(m + n)
                    @test begin
                        res = all(isapproxdefault, zip(conj.(Y₋m₋n), phase .* Ymn))
                        if !res
                            @show j, m, n, θ, ϕ
                        end
                        res
                    end
                end
            end
        end
    end
end

@testset "Coordinate inversion" begin
    jmax = 4
    @testset "Irreducible" begin
        for θ in LinRange(0, pi, 10), ϕ in LinRange(0, 2pi, 10)
            cache!(S, θ, ϕ, jmax)
            for j in 0:jmax, m in -j:j
                for B in [Cartesian(), SphericalCovariant()]
                    Yjm = vshbasis(Irreducible(), B, j, m, θ, ϕ, S)
                    Yjmrot = vshbasis(Irreducible(), B, j, m, pi - θ, pi + ϕ)
                    for l in axes(Yjm, 2)
                        Yjml = Yjm[:, l] * (-1)^(l+1)
                        Yjmlrot = -Yjmrot[:, l]
                        @test begin
                            res = all(isapproxdefault, zip(Yjml, Yjmlrot))
                            if !res
                                @show j, m, θ, ϕ
                            end
                            res
                        end
                    end
                end
            end
        end
    end
    @testset "Hansen" begin
        for θ in LinRange(0, pi, 10), ϕ in LinRange(0, 2pi, 10)
            cache!(S, θ, ϕ, jmax)
            for j in 0:4, m in -j:j
                for B in [Cartesian(), SphericalCovariant()]
                    Yjm = vshbasis(Hansen(), B, j, m, θ, ϕ, S)
                    Yjmrot = vshbasis(Hansen(), B, j, m, pi - θ, pi + ϕ)
                    for λ in axes(Yjm, 2)
                        Yjml = Yjm[:, λ] * (-1)^(j + λ +1)
                        Yjmlrot = -Yjmrot[:, λ]
                        @test all(isapproxdefault, zip(Yjml, Yjmlrot))
                    end
                end
            end
        end
    end
end

@testset "Coordinate rotation" begin
    lmax = 3
    modes = ML(0:lmax)

    S2 = VectorSphericalHarmonics.cache(0, 0, lmax);
    S′ = VectorSphericalHarmonics.cache(0, 0, lmax);
    Dvec = OffsetArray([zeros(ComplexF64, 2l+1, 2l+1) for l in 0:lmax], 0:lmax);

    @testset "arbitrary point" begin
        cartvec(θ, ϕ) = SVector{3}(sin(θ)cos(ϕ), sin(θ)sin(ϕ), cos(θ))
        function polcoords(n)
            nx, ny, nz = normalize(n)
            θ = acos(nz)
            if abs(nz) == 1
                # ϕ is not defined uniquely at the poles
                return promote(θ, zero(θ))
            else
                ϕ = atan(ny, nx)
                return promote(θ, ϕ)
            end
        end

        @testset "genspharm" begin
            for θ in LinRange(0, pi, 6), ϕ in LinRange(0, 2pi, 6)
                VectorSphericalHarmonics.cache!(S2, θ, ϕ);
                Y = genspharm(modes, θ, ϕ, S2);
                n = cartvec(θ, ϕ)
                Un = basisconversionmatrix(Cartesian(), HelicityCovariant(), θ, ϕ);

                for θ′ in LinRange(0, pi, 6), ϕ′ in LinRange(0, 2pi, 6)
                    VectorSphericalHarmonics.cache!(S′, θ′, ϕ′);
                    α, β, γ = ϕ, θ-θ′, -ϕ′
                    R = RotZYZ(α, β, γ)
                    R⁻¹ = inv(R)
                    # this specific rotation satisfies the condition U′n * R' * Un' == I,
                    # so the basis rotation matrix may be left out
                    # This happens because Un for polar == RotYZ(-θ, -ϕ) (albeit with permuted rows), so
                    # U′n * R' * Un' = RotYZ(-θ′, -ϕ′) * RotZYZ(ϕ′, θ′ - θ, -ϕ) * RotZY(ϕ, θ)  ≈ I
                    # Un for HelicityCovariant is related to that for the polar matrix through a unitary
                    # transformation, so the same relation holds
                    U′n = basisconversionmatrix(Cartesian(), HelicityCovariant(), θ′, ϕ′);
                    @test U′n * R⁻¹ * Un' ≈ I
                    for l in 0:lmax
                        Dp = wignerD!(Dvec[l], l, α, β, γ);
                        D = OffsetArray(Dp, -l:l, -l:l)
                        for m in -l:l
                            Ylθϕrot_m = sum(D[m′, m] * Y[(l, m′)] for m′ in -l:l)
                            Ylmθ′ϕ′ = genspharm(l, m, θ′, ϕ′, S′)
                            @test isapprox(Ylmθ′ϕ′, Ylθϕrot_m, atol = 1e-13, rtol = 1e-8)
                        end
                    end
                end

                # obtain the second point from an arbitrary rotation
                # In general we need to retain the basis conversion matrix to compute the components correctly
                for α in LinRange(0, 2pi, 6), β in LinRange(0, pi, 6), γ in LinRange(0, 2pi, 6)
                    # the rotation that transforms between frames S′ = RS
                    R = RotZYZ(α, β, γ)
                    R⁻¹ = inv(R)
                    n′ = R⁻¹ * n
                    θ′, ϕ′ = polcoords(n′)
                    VectorSphericalHarmonics.cache!(S′, θ′, ϕ′);
                    U′n = basisconversionmatrix(Cartesian(), HelicityCovariant(), θ′, ϕ′);
                    M⁻¹ = U′n * R⁻¹ * Un'
                    @test M⁻¹ ≈ Diagonal(M⁻¹)
                    for l in 0:lmax
                        Dp = wignerD!(Dvec[l], l, α, β, γ);
                        D = OffsetArray(Dp, -l:l, -l:l)
                        for m in -l:l
                            Ylθϕrot_m = M⁻¹ * sum(D[m′, m] * parent(Y[(l, m′)]) for m′ in -l:l)
                            Ylmθ′ϕ′ = parent(genspharm(l, m, θ′, ϕ′, S′))
                            @test isapprox(Ylmθ′ϕ′, Ylθϕrot_m, atol = 1e-13, rtol = 1e-8)
                        end
                    end
                end
            end
        end

        @testset "VSH" begin
            for θ in LinRange(0, pi, 6), ϕ in LinRange(0, 2pi, 6)
                n = cartvec(θ, ϕ)
                VectorSphericalHarmonics.cache!(S2, θ, ϕ);

                for YT in [Irreducible(), Hansen(), PB()], B in [Cartesian(), SphericalCovariant(), HelicityCovariant(), Polar()]
                    Un = basisconversionmatrix(Cartesian(), B, θ, ϕ)
                    Y = vshbasis(YT, B, modes, θ, ϕ, S2);

                    # We evaluate the rotation of the dot product between VSH and a vector field
                    # In this case we choose a constant vector field directed along x̂ at each point,
                    # where x̂ is directed along the x axis of the frame S (which differs from that of S′ in general)
                    # We evaluate the components in the basis B at both points
                    x_n = Un * SVector{3}(1,0,0)

                    if B ∈ (Polar(), HelicityCovariant())
                        for θ′ in LinRange(0, pi, 6), ϕ′ in LinRange(0, 2pi, 6)
                            VectorSphericalHarmonics.cache!(S′, θ′, ϕ′);
                            α, β, γ = ϕ, θ-θ′, -ϕ′
                            R = RotZYZ(α, β, γ)
                            # this specific rotation satisfies the condition U′n * R' * Un' == I for the
                            # Polar and HelicityCovariant bases, so the basis rotation matrix may be left out
                            # This happens because Un for polar == RotYZ(-θ, -ϕ) (albeit with permuted rows), so
                            # U′n * R' * Un' = RotYZ(-θ′, -ϕ′) * RotZYZ(ϕ′, θ′ - θ, -ϕ) * RotZY(ϕ, θ) ≈ I
                            # Un for HelicityCovariant is related to that for the polar matrix through a unitary
                            # transformation, so the same relation holds
                            U′n = basisconversionmatrix(Cartesian(), B, θ′, ϕ′);
                            if B === Polar()
                                # We go from (r,θ,ϕ) to (θ,ϕ,r)
                                @test Un[[2,3,1],:] ≈ RotYZ(-θ, -ϕ)
                                @test U′n[[2,3,1],:] ≈ RotYZ(-θ′, -ϕ′)
                            end
                            @test U′n * R' * Un' ≈ I
                            for l in 0:lmax
                                Dp = wignerD!(Dvec[l], l, α, β, γ);
                                D = OffsetArray(Dp, -l:l, -l:l)
                                for m in -l:l
                                    Ylθϕrot_m = sum(D[m′, m] * Y[(l, m′)] for m′ in -l:l)
                                    Ylmθ′ϕ′ = vshbasis(YT, B, l, m, θ′, ϕ′, S′)
                                    @test isapprox(Ylmθ′ϕ′, Ylθϕrot_m, atol = 1e-13, rtol = 1e-8)
                                end
                            end
                        end
                    end

                    # obtain the second point from an arbitrary rotation
                    # In general we need to retain the basis conversion matrix to compute the components correctly
                    for α in LinRange(0, 2pi, 6), β in LinRange(0, pi, 6), γ in LinRange(0, 2pi, 6)
                        # the rotation that transforms between frames S′ = RS
                        R = RotZYZ(α, β, γ)
                        R⁻¹ = inv(R)
                        n′ = R⁻¹ * n
                        θ′, ϕ′ = polcoords(n′)
                        VectorSphericalHarmonics.cache!(S′, θ′, ϕ′);
                        U′n = basisconversionmatrix(Cartesian(), B, θ′, ϕ′);

                        M⁻¹ = (U′n * R⁻¹ * Un')
                        if B === HelicityCovariant()
                            @test M⁻¹ ≈ Diagonal(M⁻¹)
                        end
                        M = inv(M⁻¹)

                        x_n′ = M⁻¹ * x_n

                        for l in 0:lmax
                            Dp = wignerD!(Dvec[l], l, α, β, γ);
                            D = OffsetArray(Dp, -l:l, -l:l)
                            for m in -l:l
                                Ylθϕrot_m = M⁻¹ * sum(D[m′, m] * parent(Y[(l, m′)]) for m′ in -l:l)
                                Ylmθ′ϕ′ = parent(vshbasis(YT, B, l, m, θ′, ϕ′, S′))
                                @test isapprox(Ylmθ′ϕ′, Ylθϕrot_m, atol = 1e-13, rtol = 1e-8)
                                # (M⁻¹ x)⋅Ylm′(n′) = (M⁻¹ x)⋅M⁻¹∑Dlmm′ Ylm(n) = x⋅∑Dlmm′ Ylm(n)
                                Ylθϕrot_m = x_n' * sum(D[m′, m] * parent(Y[(l, m′)]) for m′ in -l:l)
                                Ylmθ′ϕ′ = x_n′' * parent(vshbasis(YT, B, l, m, θ′, ϕ′, S′))
                                @test isapprox(Ylmθ′ϕ′, Ylθϕrot_m, atol = 1e-13, rtol = 1e-8)
                                # x.Ylm′(n′) = x.M⁻¹∑Dlmm′ Ylm(n) = (Mx)⋅∑Dlmm′ Ylm(n)
                                Ylθϕrot_m = (M*x_n)' * sum(D[m′, m] * parent(Y[(l, m′)]) for m′ in -l:l)
                                Ylmθ′ϕ′ = x_n' * parent(vshbasis(YT, B, l, m, θ′, ϕ′, S′))
                                @test isapprox(Ylmθ′ϕ′, Ylθϕrot_m, atol = 1e-13, rtol = 1e-8)
                            end
                        end
                    end
                end
            end
        end
    end
end

@testset "VSHcache" begin
    modes = LM(0:1, 0:1)
    lmax = 2
    θ1, ϕ1 = pi/3, pi/2
    θ2, ϕ2 = pi/4, pi/6
    @testset "vshbasis!" begin
        V = @inferred VectorSphericalHarmonics.VSHCache(Float64, Irreducible(), Polar(), θ1, ϕ1, modes);
        M = vshbasis(Irreducible(), Polar(), modes, θ1, ϕ1)
        @test getY(V) == M
        @test eltypeY(V) == eltype(getY(V))

        for θ in Any[θ2, NorthPole(), SouthPole()]
            vshbasis!(V, Irreducible(), Polar(), θ, ϕ2)
            M = vshbasis(Irreducible(), Polar(), modes, θ, ϕ2)
            @test getY(V) == M
        end
    end
    @testset "genspharm!" begin
        V = @inferred VectorSphericalHarmonics.VSHCache(Float64, θ1, ϕ1, modes);
        M = genspharm(modes, θ1, ϕ1)
        @test getY(V) == M

        for θ in Any[θ2, NorthPole(), SouthPole()]
            genspharm!(V, θ, ϕ2)
            M = genspharm(modes, θ, ϕ2)
            @test getY(V) == M
        end
    end
    @testset "uninitialized" begin
        V = VectorSphericalHarmonics.VSHCache(PB(), Cartesian(), LM(0:1))
        @test eltype(VectorSphericalHarmonics.eltypeY(V)) == ComplexF64
        @test all(iszero, VectorSphericalHarmonics.getY(V))

        V = VectorSphericalHarmonics.VSHCache(LM(0:1))
        @test eltype(VectorSphericalHarmonics.eltypeY(V)) == ComplexF64
        @test all(iszero, VectorSphericalHarmonics.getY(V))
    end
end

@testset "broadcast tags" begin
    function components_of_x̂(Basis, (θ,ϕ))
        B = VectorSphericalHarmonics.basisconversionmatrix(Cartesian(), Basis, θ, ϕ)
        B * SVector{3}(1,0,0)
    end
    for B in [Cartesian(), Polar(), SphericalCovariant(), HelicityCovariant()]
        x1, x2 = components_of_x̂.(B, ((pi/2, 0), (pi/2, pi/4)))
        @test x1 == components_of_x̂(B, (pi/2, 0))
        @test x2 == components_of_x̂(B, (pi/2, pi/4))
    end
    V = vshbasis.(Hansen(), Polar(), 1:2, -1, pi/2, 0)
    @test V[1] == vshbasis(Hansen(), Polar(), 1, -1, pi/2, 0)
    @test V[2] == vshbasis(Hansen(), Polar(), 2, -1, pi/2, 0)
    V = vshbasis.(Hansen(), Polar(), 1, -1, (pi/2, pi/4), 0)
    @test V[1] == vshbasis(Hansen(), Polar(), 1, -1, pi/2, 0)
    @test V[2] == vshbasis(Hansen(), Polar(), 1, -1, pi/4, 0)
end


# # Term sparsity

#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/generated/Sparsity/term_sparsity.ipynb)
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/generated/Sparsity/term_sparsity.ipynb)
# **Adapted from**: Example 3.5 of [WML20b]
#
# [WML20a] Wang, Jie, Victor Magron, and Jean-Bernard Lasserre.
# *TSSOS: A Moment-SOS hierarchy that exploits term sparsity*.
# arXiv preprint arXiv:1912.08899 (2020).
#
# [WML20b] Wang, Jie, Victor Magron, and Jean-Bernard Lasserre.
# *Chordal-TSSOS: a moment-SOS hierarchy that exploits term sparsity with chordal extension*.
# arXiv preprint arXiv:2003.03210 (2020).

using Test #src
using DynamicPolynomials
@polyvar x[1:3]

# We would like to find the minimum value of the polynomial

poly = x[1]^2 - 2x[1]*x[2] + 3x[2]^2 - 2x[1]^2*x[2] + 2x[1]^2*x[2]^2 - 2x[2]*x[3] + 6x[3]^2 + 18x[2]^2*x[3] - 54x[2]*x[3]^2 + 142x[2]^2*x[3]^2

# The minimum value of the polynomial can be found to be zero.

import CSDP
solver = CSDP.Optimizer
using SumOfSquares
function sos_min(sparsity)
    model = Model(solver)
    @variable(model, t)
    @objective(model, Max, t)
    con_ref = @constraint(model, poly - t in SOSCone(), sparsity = sparsity)
    optimize!(model)
    return value(t), moment_matrix(con_ref)
end

bound, ν = sos_min(NoSparsity())
@test bound ≈ 0 atol=1e-6 #src
bound

# We find the corresponding minimizer `(0, 0, 0)` by matching the moments
# of the moment matrix with a dirac measure centered at this minimizer.

extractatoms(ν, 1e-6)

# We can see below that the basis contained 6 monomials hence we needed to use 6x6 PSD matrix variables.

@test ν.basis.monomials == [x[1]*x[2], x[2]*x[3], x[1], x[2], x[3], 1] #src
ν.basis

# Using the monomial/term sparsity method of [WML20a] based on cluster completion, we find the same bound.

bound, ν = sos_min(MonomialSparsity())
@test bound ≈ 0 atol=1e-6 #src
bound

# Which is not suprising as no sparsity reduction could be performed.

@test length(ν.sub_moment_matrices) == 1 #src
@test ν.sub_moment_matrices[1].basis.monomials == [x[1]*x[2], x[2]*x[3], x[1], x[2], x[3], 1] #src
[sub.basis for sub in ν.sub_moment_matrices]

# Using the monomial/term sparsity method of [WML20b] based on chordal completion, the lower bound is smaller than 0.

bound, ν = sos_min(MonomialSparsity(ChordalCompletion()))
@test bound ≈ -0.00355 rtol=1e-3 #src
bound

# However, this bound was obtained with an SDP with 4 matrices of size 3x3.

@test length(ν.sub_moment_matrices) == 4                                  #src
@test ν.sub_moment_matrices[1].basis.monomials == [x[2]*x[3], x[2], x[3]] #src
@test ν.sub_moment_matrices[2].basis.monomials == [x[2]*x[3], x[2], 1]    #src
@test ν.sub_moment_matrices[3].basis.monomials == [x[1], x[2], 1]         #src
@test ν.sub_moment_matrices[4].basis.monomials == [x[1]*x[2], x[1], 1]    #src
[sub.basis for sub in ν.sub_moment_matrices]


@testset "calc_position func" begin
    time_vec = [0.0, 2.0]
    u_vec = [1.0, 1.0]
    v_vec = [0.0, 0.0]
    r_vec = [0.0, 0.0]
    x, y, ψ = calc_position(time_vec, u_vec, v_vec, r_vec, x0 = 1.0, y0 = 1.0, ψ0 = 0.0)
    @test x[2] ≈ 3.0
    @test y[2] ≈ 1.0
    @test ψ[2] ≈ 0.0
end


## Root finding for Chebyshev expansions
#
#  Contains code that is based in part on Chebfun v5's chebfun/@chebteck/roots.m,
# which is distributed with the following license:

# Copyright (c) 2015, The Chancellor, Masters and Scholars of the University
# of Oxford, and the Chebfun Developers. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#     * Redistributions of source code must retain the above copyright
#       notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#     * Neither the name of the University of Oxford nor the names of its
#       contributors may be used to endorse or promote products derived from
#       this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

function complexroots(f::Fun{C}) where C<:Chebyshev
    if ncoefficients(f)==0 || (ncoefficients(f)==1 && isapprox(f.coefficients[1],0))
        throw(ArgumentError("Tried to take roots of a zero function."))
    elseif ncoefficients(f)==1
        ComplexF64[]
    elseif ncoefficients(f)==2
        ComplexF64[-f.coefficients[1]/f.coefficients[2]]
    else
        fromcanonical.(f,colleague_eigvals(f.coefficients))
    end
end

#function roots(f::Fun)
#    irts=map(real,filter!(x->abs(x)<=1.+10eps(),filter!#(isreal,complexroots(f.coefficients))))
#
#    map!(x->x>1. ? 1. : x,irts)
#    map!(x->x<-1. ? -1. : x,irts)
#
#    if length(irts)==0
#        Float64[]
#    else
#        fromcanonical(f,irts)
#    end
#end

function roots(f::Fun)
    f2=Fun(f,domain(f)) # default is to convert to Chebyshev/Fourier
    if space(f2)==space(f)
        error("roots not implemented for "*string(typeof(f)))
    else
        roots(f2)
    end
end

roots(f::Fun{<:Chebyshev}) = fromcanonical.(f,roots(setcanonicaldomain(f)))

for (BF,FF) in ((BigFloat,Float64),(Complex{BigFloat},ComplexF64))
    @eval function roots( f::Fun{C,$BF} ) where C<:Chebyshev
    # FIND THE ROOTS OF AN IFUN.

        d = domain(f)
        c = f.coefficients
        vscale = maximum(abs,values(f))
        if vscale == 0
            throw(ArgumentError("Tried to take roots of a zero function."))
        end

        hscale = maximum( [abs(leftendpoint(d)), abs(rightendpoint(d))] )
        htol = eps(2000.)*max(hscale, 1)  # TODO: choose tolerance better

        # calculate Flaot64 roots
        r = Array{$BF}(rootsunit_coeffs(convert(Vector{$FF},c./vscale), Float64(htol)))
        # Map roots from [-1,1] to domain of f:
        rts = fromcanonical.(Ref(d),r)
        fp = differentiate(f)

        # do Newton 3 time
        for _ = 1:3
            rts .-=extrapolate.(Ref(f),rts)./extrapolate.(Ref(fp),rts)
        end

        return rts
    end
end

function roots( f::Fun{C,TT} ) where {C<:Chebyshev,TT<:Union{Float64,ComplexF64}}
# FIND THE ROOTS OF AN IFUN.
    if iszero(f)
        throw(ArgumentError("Tried to take roots of a zero function."))
    end

    d = domain(f)
    c = f.coefficients
    vscale = maximum(abs,values(f))

    hscale = max(abs(leftendpoint(d)), abs(rightendpoint(d)) )
    htol = eps(2000.)*max(hscale, 1)  # TODO: choose tolerance better

    cvscale=c./vscale
    r = rootsunit_coeffs(cvscale, Float64(htol))

    # Check endpoints, as these are prone to inaccuracy
    # which can be deadly.
    if (isempty(r) || !isapprox(last(r),1.)) && abs(sum(cvscale)) < htol
        push!(r,1.)
    end
    if (isempty(r) || !isapprox(first(r),-1.)) && abs(alternatingsum(cvscale)) < htol
        insert!(r,1,-1.)
    end

    # Map roots from [-1,1] to domain of f:
    return fromcanonical.(Ref(d),r)
end

function colleague_matrix( c::Vector{T} ) where T<:Number
#TODO: This is command isn't typed correctly
# COMPUTE THE ROOTS OF A LOW DEGREE POLYNOMIAL BY USING THE COLLEAGUE MATRIX:
    c = chop(c,0) # prune exact zeros
    n = length(c) - 1
    A=zeros(T,n,n)

    for k=1:n-1
        A[k+1,k]=0.5
        A[k,k+1]=0.5
    end
    for k=1:n
        A[1,end-k+1]-=0.5*c[k]/c[end]
    end
    A[n,n-1]=one(T)

    # TODO: can we speed things up because A is lower-Hessenberg
    # Standard colleague matrix (See [Good, 1961]):
    A
end

function colleague_balance!(M)
    n=size(M,1)
    if n<2
        return M
    end

    d=1+abs(M[1,2])
    M[1,2]/=d;M[2,1]*=d

    for k=3:n
        M[k-1,k]*=d;M[k,k-1]/=d
        d+=abs(M[1,k])
        M[1,k]/=d;M[k-1,k]/=d;M[k,k-1]*=d
    end
    M
end

# colleague_eigvals( c::Vector )=hesseneigvals(colleague_balance!(colleague_matrix(c)))

colleague_eigvals( c::Vector )=eigvals(colleague_matrix(c))

function PruneOptions( r, htol::Float64 )
# ONLY KEEP ROOTS IN THE INTERVAL

    # Remove dangling imaginary parts:
    r = real( r[ abs.(imag.(r)) .< htol ] )
    # Keep roots inside [-1 1]:
    r = sort( r[ abs.(r) .<= 1+htol ] )
    # Put roots near ends onto the domain:
    r = min.( max.( r, -1 ), 1)

    # Return the pruned roots:
    return r
end

rootsunit_coeffs(c::Vector{T}, htol::Float64) where {T<:Number} =
    rootsunit_coeffs(c,htol,ClenshawPlan(T,Chebyshev(),length(c),length(c)))
function rootsunit_coeffs(c::Vector{T}, htol::Float64,clplan::ClenshawPlan{S,T}) where {S,T<:Number}
# Computes the roots of the polynomial given by the coefficients c on the unit interval.

    # If recursive subdivision is used, then subdivide [-1,1] into [-1,splitPoint] and [splitPoint,1].
    splitPoint = -0.004849834917525

    # Simplify the coefficients by chopping off the tail:
    nrmc=norm(c, 1)
    @assert nrmc > 0  # There's an error if we make it this far
    c = chop(c,eps()*nrmc)
    n = length(c)

    if n == 0

        # EMPTY FUNCTION
        r = Float64[]

    elseif n == 1

        # CONSTANT FUNCTION
        r = ( c[1] == 0.0 ) ? Float64[0.0] : Float64[]

    elseif n == 2

        # LINEAR POLYNOMIAL
        r1 = -c[1]/c[2];
        r = ( (abs(imag(r1))>htol) | (abs(real(r1))>(1+htol)) ) ? Float64[] : Float64[max(min(real(r1),1),-1)]

    elseif n <= 70

        # COLLEAGUE MATRIX
        # The recursion subdividing below will keep going until we have a piecewise polynomial
        # of degree at most 50 and we likely end up here for each piece.

        # Adjust the coefficients for the colleague matrix
        # Prune roots depending on preferences:
        r = PruneOptions( colleague_eigvals(c), htol )::Vector{Float64}
    else

        #  RECURSIVE SUBDIVISION:
        # If n > 50, then split the interval [-1,1] into [-1,splitPoint], [splitPoint,1].
        # Find the roots of the polynomial on each piece and then concatenate. Recurse if necessary.

        # Evaluate the polynomial at Chebyshev grids on both intervals:
        #(clenshaw! overwrites points, which only makes sence if c is real)

        v1 = isa(c,Vector{Float64}) ? clenshaw!( c, points(-1..splitPoint,n),clplan) : clenshaw( c, points(-1..splitPoint,n),clplan)
        v2 = isa(c,Vector{Float64}) ? clenshaw!( c, points(splitPoint..1 ,n),clplan) : clenshaw( c, points(splitPoint..1 ,n),clplan)

        # Recurse (and map roots back to original interval):
        p = plan_chebyshevtransform( v1 )
        r = [ (splitPoint - 1)/2 .+ (splitPoint + 1)/2 .* rootsunit_coeffs( p*v1, 2*htol,clplan) ;
                 (splitPoint + 1)/2 .+ (1 - splitPoint)/2 .* rootsunit_coeffs( p*v2, 2*htol,clplan) ]

    end

    # Return the roots:
    r
end

extremal_args(f::Fun{S}) where {S<:ContinuousSpace} = cat(1,[extremal_args(fp) for fp in components(f)]..., dims=1)

for op in (:(maximum),:(minimum))
    @eval begin

        $op(f::Fun{ContinuousSpace{T,R},T}) where {T<:Real,R<:Real} =
            $op(map($op,components(f)))
        $op(::typeof(abs), f::Fun{ContinuousSpace{T,R},T}) where {T<:Real,R<:Real} =
            $op(abs, map(g -> $op(abs, g),components(f)))
    end
end

extrema(f::Fun{ContinuousSpace{T,R},T}) where {T<:Real,R<:Real} =
    mapreduce(extrema,(x,y)->extrema([x...;y...]),components(f))

function roots(f::Fun{P}) where P<:ContinuousSpace
    rts=mapreduce(roots,vcat,components(f))
    k=1
    while k < length(rts)
        if isapprox(rts[k],rts[k+1])
            rts=rts[[1:k;k+2:end]]
        else
            k+=1
        end
    end

    rts
end


using Clustering
using GaussianMixtures
using LinearAlgebra

mutable struct MOE{X,Y,L,U,S,K,M,V,W} <: AbstractSurrogate
    x::X
    y::Y
    lb::L
    ub::U
    local_surr::S
    k::K
    means::M
    varcov::V
    weights::W
end

#Radial structure:
function RadialBasisStructure(;radial_function,scale_factor,sparse)
    return (name = "RadialBasis", radial_function = radial_function, scale_factor = scale_factor, sparse = sparse)
end

#Kriging structure:
function KrigingStructure(;p,theta)
    return (name = "Kriging", p = p, theta = theta)
end

#Linear structure
function LinearStructure()
    return (name = "LinearSurrogate")
end

#InverseDistance structure
function InverseDistanceStructure(;p)
    return (name = "InverseDistanceSurrogate", p = p)
end

#Lobachesky structure
function LobacheskyStructure(;alpha,n,sparse)
    return (name = "LobacheskySurrogate", alpha = alpha, n = n, sparse = sparse)
end

#Neural structure
function NeuralStructure(;model,loss,opt,n_echos)
    return (name ="NeuralSurrogate", model = model ,loss = loss,opt = opt,n_echos = n_echos)
end

#Random forest structure
function RandomForestStructure(;num_round)
    return (name = "RandomForestSurrogate", num_round = num_round)
end

#Second order poly structure
function SecondOrderPolynomialStructure()
    return (name = "SecondOrderPolynomialSurrogate")
end

#Wendland structure
function WendlandStructure(; eps, maxiters, tol)
    return (name = "Wendland", eps = eps, maxiters = maxiters, tol = tol)
end

#Polychaos structure
function PolyChaosStructure(; op)
    return (name = "PolynomialChaosSurrogate", op = op)
end

function MOE(x,y,lb::Number,ub::Number; k::Int = 2, local_kind = [RadialBasisStructure(radial_function = linearRadial, scale_factor=1.0,sparse = false),RadialBasisStructure(radial_function = cubicRadial, scale_factor=1.0, sparse = false)])
    if k != length(local_kind)
        throw("Number of mixtures = $k is not equal to length of local surrogates")
    end
    n = length(x)
    # find weight, mean and variance for each mixture
    # For GaussianMixtures I need nxd matrix
    X_G = reshape(x,(n,1))
    moe_gmm = GaussianMixtures.GMM(k,X_G)
    weights = GaussianMixtures.weights(moe_gmm)
    means = GaussianMixtures.means(moe_gmm)
    variances = moe_gmm.Σ

    #cluster the points
    #For clustering I need dxn matrix
    X_C = reshape(x,(1,n))
    KNN = kmeans(X_C, k)
    x_c = [ Array{eltype(x)}(undef,0) for i = 1:k]
    y_c = [ Array{eltype(y)}(undef,0) for i = 1:k]
    a = assignments(KNN)
    @inbounds for i = 1:n
        pos = a[i]
        append!(x_c[pos],x[i])
        append!(y_c[pos],y[i])
    end

    local_surr = Dict()
    for i = 1:k
        if local_kind[i][1] == "RadialBasis"
            #fit and append to local_surr
            my_local_i = RadialBasis(x_c[i],y_c[i],lb,ub,rad = local_kind[i].radial_function, scale_factor = local_kind[i].scale_factor, sparse = local_kind[i].sparse)
            local_surr[i] = my_local_i

        elseif local_kind[i][1] == "Kriging"
            my_local_i = Kriging(x_c[i], y_c[i],lb,ub, p = local_kind[i].p, theta = local_kind[i].theta)
            local_surr[i] = my_local_i

        elseif local_kind[i] == "LinearSurrogate"
            my_local_i = LinearSurrogate(x_c[i], y_c[i],lb,ub)
            local_surr[i] = my_local_i

        elseif local_kind[i][1] == "InverseDistanceSurrogate"
            my_local_i = InverseDistanceSurrogate(x_c[i], y_c[i],lb,ub, local_kind[i].p)
            local_surr[i] = my_local_i

        elseif local_kind[i][1] == "LobacheskySurrogate"
            my_local_i = LobacheskySurrogate(x_c[i], y_c[i],lb,ub,alpha = local_kind[i].alpha , n = local_kind[i].n, sparse = local_kind[i].sparse)
            local_surr[i] = my_local_i

        elseif local_kind[i][1] == "NeuralSurrogate"
            my_local_i = NeuralSurrogate(x_c[i], y_c[i],lb,ub, model = local_kind[i].model , loss = local_kind[i].loss ,opt = local_kind[i].opt ,n_echos = local_kind[i].n_echos)
            local_surr[i] = my_local_i

        elseif local_kind[i][1] == "RandomForestSurrogate"
            my_local_i = RandomForestSurrogate(x_c[i], y_c[i],lb,ub, num_round = local_kind[i].num_round)
            local_surr[i] = my_local_i

        elseif local_kind[i] == "SecondOrderPolynomialSurrogate"
            my_local_i = SecondOrderPolynomialSurrogate(x_c[i], y_c[i],lb,ub)
            local_surr[i] = my_local_i

        elseif local_kind[i][1] == "Wendland"
            my_local_i = Wendand(x_c[i], y_c[i],lb,ub, eps = local_kind[i].eps, maxiters = local_kind[i].maxiters, tol = local_kind[i].tol)
            local_surr[i] = my_local_i
        elseif local_kind[i][1] == "PolynomialChaosSurrogate"
            my_local_i = PolynomialChaosSurrogate(x,y,lb,ub, op = local_kind[i].op)
            local_surr[i] = my_local_i

        else
            throw("A surrogate with name provided does not exist or is not currently supported with MOE.")
        end
    end
    return MOE(x,y,lb,ub,local_surr,k,means,variances,weights)
end

function MOE(x,y,lb,ub; k::Int = 2,
            local_kind = [RadialBasisStructure(radial_function = linearRadial, scale_factor=1.0, sparse = false),RadialBasisStructure(radial_function = cubicRadial, scale_factor=1.0, sparse = false)])

    n = length(x)
    d = length(lb)
    #GMM parameters:
    X_G = collect(reshape(collect(Base.Iterators.flatten(x)), (d,n))')
    my_gmm = GMM(k,X_G,kind = :full)
    weights = my_gmm.w
    means = my_gmm.μ
    varcov = my_gmm.Σ

    #cluster the points
    X_C = collect(reshape(collect(Base.Iterators.flatten(x)), (d,n)))
    KNN = kmeans(X_C, k)
    x_c = [ Array{eltype(x)}(undef,0) for i = 1:k]
    y_c = [ Array{eltype(y)}(undef,0) for i = 1:k]
    a = assignments(KNN)
    @inbounds for i = 1:n
        pos = a[i]
        x_c[pos] = vcat(x_c[pos],x[i])
        append!(y_c[pos],y[i])
    end

    local_surr = Dict()
    for i = 1:k
        if local_kind[i][1] == "RadialBasis"
            #fit and append to local_surr
            my_local_i = RadialBasis(x_c[i],y_c[i],lb,ub,rad = local_kind[i].radial_function, scale_factor = local_kind[i].scale_factor, sparse = local_kind[i].sparse)
            local_surr[i] = my_local_i

        elseif local_kind[i][1] == "Kriging"
            my_local_i = Kriging(x_c[i], y_c[i],lb,ub, p = local_kind[i].p, theta = local_kind[i].theta)
            local_surr[i] = my_local_i

        elseif local_kind[i] == "LinearSurrogate"
            my_local_i = LinearSurrogate(x_c[i], y_c[i],lb,ub)
            local_surr[i] = my_local_i

        elseif local_kind[i][1] == "InverseDistanceSurrogate"
            my_local_i = InverseDistanceSurrogate(x_c[i], y_c[i],lb,ub, local_kind[i].p)
            local_surr[i] = my_local_i

        elseif local_kind[i][1] == "LobacheskySurrogate"
            my_local_i = LobacheskySurrogate(x_c[i], y_c[i],lb,ub,alpha = local_kind[i].alpha , n = local_kind[i].n, sparse = local_kind[i].sparse)
            local_surr[i] = my_local_i

        elseif local_kind[i][1] == "NeuralSurrogate"
            my_local_i = NeuralSurrogate(x_c[i], y_c[i],lb,ub, model = local_kind[i].model , loss = local_kind[i].loss ,opt = local_kind[i].opt ,n_echos = local_kind[i].n_echos)
            local_surr[i] = my_local_i

        elseif local_kind[i][1] == "RandomForestSurrogate"
            my_local_i = RandomForestSurrogate(x_c[i], y_c[i],lb,ub, num_round = local_kind[i].num_round)
            local_surr[i] = my_local_i

        elseif local_kind[i] == "SecondOrderPolynomialSurrogate"
            my_local_i = SecondOrderPolynomialSurrogate(x_c[i], y_c[i],lb,ub)
            local_surr[i] = my_local_i

        elseif local_kind[i][1] == "Wendland"
            my_local_i = Wendand(x_c[i], y_c[i],lb,ub, eps = local_kind[i].eps, maxiters = local_kind[i].maxiters, tol = local_kind[i].tol)
            local_surr[i] = my_local_i

        elseif local_kind[i][1] == "PolynomialChaosSurrogate"
            my_local_i = PolynomialChaosSurrogate(x,y,lb,ub, op = local_kind[i].op)
            local_surr[i] = my_local_i
        else
            throw("A surrogate with name "* local_kind[i][1] *" does not exist or is not currently supported with MOE.")
        end
    end
    return MOE(x,y,lb,ub,local_surr,k,means,varcov,weights)
end

function _prob_x_in_i(x::Number,i,mu,varcov,alpha,k)
    num = (1/sqrt(varcov[i]))*alpha[i]*exp(-0.5(x-mu[i])*(1/varcov[i])*(x-mu[i]))
    den = sum([(1/sqrt(varcov[j]))*alpha[j]*exp(-0.5(x-mu[j])*(1/varcov[j])*(x-mu[j]))  for j = 1:k])
    return num/den
end

function _prob_x_in_i(x,i,mu,varcov,alpha,k)
    num = (1/sqrt(det(varcov[i])))*alpha[i]*exp(-0.5*(x .- mu[i,:])'*(inv(varcov[i]))*(x .- mu[i,:]))
    den = sum([(1/sqrt(det(varcov[j])))*alpha[j]*exp(-0.5*(x .- mu[j,:])'*(inv(varcov[j]))*(x .- mu[j,:])) for j = 1:k])
    return num/den
end

function (moe::MOE)(val)
    return prod([moe.local_surr[i](val)*_prob_x_in_i(val,i,moe.means,moe.varcov,moe.weights,moe.k) for i = 1:moe.k])
end

function add_point!(my_moe::MOE,x_new,y_new)
    if length(my_moe.x[1]) == 1
        #1D
        my_moe.x = vcat(my_moe.x,x_new)
        my_moe.y = vcat(my_moe.y,y_new)
        n = length(my_moe.x)

        #New mixture parameters
        X_G = reshape(my_moe.x,(n,1))
        moe_gmm = GaussianMixtures.GMM(my_moe.k,X_G)
        my_moe.weights = GaussianMixtures.weights(moe_gmm)
        my_moe.means = GaussianMixtures.means(moe_gmm)
        my_moe.varcov = moe_gmm.Σ

        #Find cluster of new point(s):
        n_added = length(x_new)
        X_C = reshape(my_moe.x,(1,n))
        KNN = kmeans(X_C, my_moe.k)
        a = assignments(KNN)
        #Recalculate only relevant surrogates
        for i = 1:n_added
            pos = a[n-n_added+i]
            add_point!(my_moe.local_surr[i],my_moe.x[n-n_added+i],my_moe.y[n-n_added+i])
        end
    else
        #ND
        my_moe.x = vcat(my_moe.x,x_new)
        my_moe.y = vcat(my_moe.y,y_new)
        n = length(my_moe.x)
        d = length(my_moe.lb)
        #New mixture parameters
        X_G = collect(reshape(collect(Base.Iterators.flatten(my_moe.x)), (d,n))')
        my_gmm = GMM(my_moe.k,X_G,kind = :full)
        my_moe.weights = my_gmm.w
        my_moe.means = my_gmm.μ
        my_moe.varcov = my_gmm.Σ

        #cluster the points
        X_C = collect(reshape(collect(Base.Iterators.flatten(my_moe.x)), (d,n)))
        KNN = kmeans(X_C, my_moe.k)
        a = assignments(KNN)
        n_added = length(x_new)
        for i = 1:n_added
            pos = a[n-n_added+i]
            add_point!(my_moe.local_surr[i],my_moe.x[n-n_added+i],my_moe.y[n-n_added+i])
        end
    end
    nothing
end


# ## Data

import CSV
import Random
import TuringModels

using DataFrames
using RData
using LinearAlgebra

Random.seed!(1)

data_dir = joinpath(TuringModels.project_root, "data")
kline_path = joinpath(data_dir, "Kline2.csv")
df = CSV.read(kline_path, DataFrame; delim=';')

dmat_path = joinpath(data_dir, "islandsDistMatrix.rda")
dmat = load(dmat_path)["islandsDistMatrix"]

df.society = 1:10;
df

# ## Model

using Turing

@model function m13_7(Dmat, society, logpop, total_tools)
    rhosq ~ truncated(Cauchy(0, 1), 0, Inf)
    etasq ~ truncated(Cauchy(0, 1), 0, Inf)
    bp ~ Normal(0, 1)
    a ~ Normal(0, 10)

    ## GPL2
    SIGMA_Dmat = etasq * exp.(-rhosq * Dmat.^2)
    SIGMA_Dmat = SIGMA_Dmat + 0.01I
    SIGMA_Dmat = (SIGMA_Dmat' + SIGMA_Dmat) / 2
    g ~ MvNormal(zeros(10), SIGMA_Dmat)

    log_lambda = a .+ g[society] .+ bp * logpop

    total_tools .~ Poisson.(exp.(log_lambda))
end

chns = sample(
    m13_7(dmat, df.society, df.logpop, df.total_tools),
    NUTS(),
    5000
)

# ## Output

# \defaultoutput{}

# ## Original output

m_13_7_rethinking = """
Inference for Stan model: 6422d8042e9cdd08dae2420ad26842f1.
    4 chains, each with iter=10000; warmup=2000; thin=1;
    post-warmup draws per chain=8000, total post-warmup draws=32000.

            mean se_mean    sd   2.5%    25%    50%    75%  97.5% n_eff Rhat
    g[1]   -0.27    0.01  0.44  -1.26  -0.50  -0.24  -0.01   0.53  3278    1
    g[2]   -0.13    0.01  0.43  -1.07  -0.34  -0.10   0.12   0.67  3144    1
    g[3]   -0.17    0.01  0.42  -1.10  -0.37  -0.14   0.07   0.59  3096    1
    g[4]    0.30    0.01  0.37  -0.51   0.12   0.30   0.50   1.02  3145    1
    g[5]    0.03    0.01  0.37  -0.77  -0.15   0.04   0.22   0.72  3055    1
    g[6]   -0.46    0.01  0.38  -1.31  -0.64  -0.42  -0.23   0.19  3306    1
    g[7]    0.10    0.01  0.36  -0.70  -0.07   0.11   0.29   0.78  3175    1
    g[8]   -0.26    0.01  0.37  -1.07  -0.43  -0.24  -0.06   0.40  3246    1
    g[9]    0.23    0.01  0.35  -0.52   0.07   0.25   0.42   0.89  3302    1
    g[10]  -0.12    0.01  0.45  -1.04  -0.36  -0.11   0.13   0.77  5359    1
    a       1.31    0.02  1.15  -0.98   0.61   1.31   2.00   3.69  4389    1
    bp      0.25    0.00  0.11   0.02   0.18   0.24   0.31   0.47  5634    1
    etasq   0.34    0.01  0.53   0.03   0.10   0.20   0.39   1.52  5643    1
    rhosq   1.52    0.09 11.82   0.03   0.16   0.39   0.96   7.96 15955    1
    lp__  925.98    0.03  2.96 919.16 924.20 926.34 928.14 930.67  7296    1
""";


@doc raw"""
    adjoint_differential_geodesic_startpoint(M,p, q, t, X)

Compute the adjoint of $D_p γ(t; p, q)[X]$.

# See also

[`differential_geodesic_startpoint`](@ref), [`adjoint_Jacobi_field`](@ref)
"""
adjoint_differential_geodesic_startpoint(M::MT, p, q, t, X) where {MT <: Manifold} = adjoint_Jacobi_field(M, p, q, t, X, βdifferential_geodesic_startpoint)

@doc raw"""
    adjoint_differential_geodesic_endpoint(M, p, q, t, X)

Compute the adjoint of $D_q γ(t; p, q)[X]$.

# See also

[`differential_geodesic_endpoint`](@ref), [`adjoint_Jacobi_field`](@ref)
"""
function adjoint_differential_geodesic_endpoint(M::MT, p, q, t, X) where {MT <: Manifold}
    return adjoint_Jacobi_field(M, q, p, 1-t, X, βdifferential_geodesic_startpoint)
end

@doc raw"""
    adjoint_differential_exp_basepoint(M, p, X, Y)

Computes the adjoint of $D_p \exp_p X[Y]$.

# See also

[`differential_exp_basepoint`](@ref), [`adjoint_Jacobi_field`](@ref)
"""
function adjoint_differential_exp_basepoint(M::MT, p, X, Y) where {MT <: Manifold}
    return adjoint_Jacobi_field(M, p, exp(M, p, X), 1., Y, βdifferential_exp_basepoint)
end

@doc raw"""
    adjoint_differential_exp_argument(M, p, X, Y)

Compute the adjoint of $D_X\exp_p X[Y]$.
Note that $X ∈  T_p(T_p\mathcal M) = T_p\mathcal M$ is still a tangent vector.

# See also

[`differential_exp_argument`](@ref), [`adjoint_Jacobi_field`](@ref)
"""
function adjoint_differential_exp_argument(M::mT, p, X, Y) where {mT <: Manifold}
    return adjoint_Jacobi_field(M, p, exp(M, p, X), 1.0, Y, βdifferential_exp_argument)
end

@doc raw"""
    adjoint_differential_log_basepoint(M, p, q, X)

computes the adjoint of $D_p log_p q[X]$.

# See also
[`differential_log_basepoint`](@ref), [`adjoint_Jacobi_field`](@ref)
"""
function adjoint_differential_log_basepoint(M::Manifold, p, q, X)
    return adjoint_Jacobi_field(M, p, q, 0., X, βdifferential_log_basepoint)
end

@doc raw"""
    adjoint_differential_log_argument(M, p, q, X)

Compute the adjoint of $D_q log_p q[X]$.

# See also
[`differential_log_argument`](@ref), [`adjoint_Jacobi_field`](@ref)
"""
function adjoint_differential_log_argument(M::Manifold, p, q, X)
    return adjoint_Jacobi_field(M, p, q, 1., X, βdifferential_log_argument)
end

@doc raw"""
    Y = adjoint_differential_forward_logs(M, p, X)

Compute the adjoint differential of [`forward_logs`](@ref) $F$ orrucirng,
in the power manifold array `p`, the differential of the function

$F_i(p) = \sum_{j ∈ \mathcal I_i} \log_{p_i} p_j$

where $i$ runs over all indices of the `PowerManifold` manifold `M` and $\mathcal I_i$
denotes the forward neighbors of $i$
Let $n$ be the number dimensions of the `PowerManifold` manifold (i.e. `length(size(x)`)).
Then the input tangent vector lies on the manifold $\mathcal M' = \mathcal M^n$.

# Input

* `M`     – a `PowerManifold` manifold
* `p`     – an array of points on a manifold
* `X`     – a tangent vector to from the n-fold power of `p`, where n is the `ndims` of `p`

# Ouput

`Y` – resulting tangent vector in $T_p\mathcal M$ representing the adjoint
  differentials of the logs.
"""
function adjoint_differential_forward_logs(M::PowerManifold{𝔽,TM,TSize,TPR}, p, X) where {𝔽,TM,TSize,TPR}
    power_size = power_dimensions(M)
    R = CartesianIndices(Tuple(power_size))
    d = length(power_size)
    maxInd = [last(R).I...] # maxInd as Array
    N = PowerManifold(M.manifold, TPR(), power_size..., d)
    Y = zero_tangent_vector(M,p)
    for i in R # iterate over all pixel
        for k in 1:d # for all direction combinations
            I = [i.I...] # array of index
            J = I .+ 1 .* (1:d .== k) #i + e_k is j
            if all( J .<= maxInd ) # is this neighbor in range?
                j = CartesianIndex{d}(J...) # neigbbor index as Cartesian Index
                # Y[i] += adjoint_differential_log_basepoint(M.manifold,p[i],p[j],X[i,k])
                set_component!(
                    M,
                    Y,
                    get_component(M,Y,i) .+ adjoint_differential_log_basepoint(
                        M.manifold,
                        get_component(M,p,i),
                        get_component(M,p,j),
                        get_component(N,X,i,k),
                    ),
                    i,
                )
                # Y[j] += adjoint_differential_log_argument(M.manifold,p[i],p[j],X[i,k])
                set_component!(
                    M,
                    Y,
                    get_component(M,Y,j) + adjoint_differential_log_basepoint(
                        M.manifold,
                        get_component(M,p,i),
                        get_component(M,p,j),
                        get_component(N,X,i,k),
                    ),
                    i,
                )
            end
        end # directions
    end # i in R
    return Y
end


# NOT CURRENTLY WORKING !!! NEED TO ADD TIME STEPPING ROUTINES TO eFEM

using eFEMpart, JLD

using LinearAlgebra # for condition number calculation

function main()
  function computeNorms(dt,i)
    start = time()
    mesh = squareMesh([-2,2.0,-1.,1.],16,2);
    Tfinal = 0.5
    varname      = "temperature"
    OperatorType = :Poisson2D     # -Delta u = f

    # Define auxiliary information
    dNodes = Dirichlet(:left,:right)
    nNodes = Neumann(:bottom,:top)
    dBCf = Dirichlet((x,y) -> 0.0)
    nBCf = Neumann((x,y) -> 0.0)

    # required technicality for grouping auxiliary information
    Nodes = [dNodes,nNodes]; bcfun = [dBCf,nBCf]

    # define problem variable (helps solver)
    prob = Problem(mesh,Nodes,bcfun,varname,OperatorType)

    param = PoissonParam(1.0)
    xic = [mesh.xy[i].x for i=1:length(mesh.xy)]
    U0  = sin.(pi*xic/2)

    U = FEMForwardEuler(prob,mesh,param,U0,Tfinal,dt)
    sol = ScalarSolution(U)

    #fn = Path("HeatValidation$(i)")
    # print problem
    #vtksave(mesh,prob,param,sol,fn)

    elapsed = time()-start

    # generate exact solution
    Uexact(x,y) = exp(-(pi/2)^2*Tfinal)*sin(pi*x/2)
    UexactArr = [Uexact.(mesh.xy[i].x,mesh.xy[i].y) for i=1:length(mesh.xy)]

    # compute difference
    err = sol.u - UexactArr

    # compute Li norm
    L1   = DomainNorm(mesh.xy,mesh.cm,err;normID="1")
    L2   = DomainNorm(mesh.xy,mesh.cm,err;normID="2")
    Linf = DomainNorm(mesh.xy,mesh.cm,err;normID="Inf")

    return L1,L2,Linf,elapsed
  end

  start = 1e-4

  dtarr    = [start,start,start/2,start/4,start/8,start/16]
  N       = length(dtarr)
  L1arr   = zeros(N)
  L2arr   = zeros(N)
  LInfarr = zeros(N)
  timearr = zeros(N)

  for i=1:N
      L1arr[i],L2arr[i],LInfarr[i],timearr[i] = computeNorms(dtarr[i],i)
      println("completed i=",i,"/",N)
  end

  L1arr = L1arr[2:end]
  L2arr = L2arr[2:end]
  LInfarr = LInfarr[2:end]
  timearr = timearr[2:end]
  dtarr = dtarr[2:end]

  run(`mkdir -p data`)

  save("data/TEMPHeat_validation.jld", "L1arr", L1arr,
                                        "L2arr", L2arr,
                                        "LInfarr",LInfarr,
                                        "timearr",timearr,
                                        "dtarr",dtarr)
  nothing
end

main()


export getSensTMatVec

function getSensTMatVec(v::Vector,m::Vector,pFor::FWIparam)

    # extract pointers
    M  				= pFor.Mesh
    omega 			= pFor.omega
    gamma 			= pFor.gamma
    Q     			= pFor.Sources
    P     			= pFor.Receivers
    Ainv			= pFor.ForwardSolver

	batchSize 		= pFor.forwardSolveBatchSize;
	select  		= pFor.sourceSelection;

	nsrc = size(Q,2); nrec = size(P,2);

	if length(select) > 0
		nsrc = length(select);
	end

	if batchSize > nsrc
		batchSize = nsrc;
	end

	numBatches 	= ceil(Int64,nsrc/batchSize);
	n = prod(M.n.+1);

	H = spzeros(ComplexF64,n,n);
	if isa(Ainv,ShiftedLaplacianMultigridSolver)
		H = GetHelmholtzOperator(M,m,omega, gamma, true,useSommerfeldBC);
		Ainv.helmParam = HelmholtzParam(M,gamma,m,omega,true,useSommerfeldBC);
		H = H + GetHelmholtzShiftOP(m, omega,Ainv.shift[1]);
		H = sparse(H');
		# H is actually shifted laplacian now...
	elseif isa(Ainv,JuliaSolver)
		H = GetHelmholtzOperator(M,m,omega, gamma, true,useSommerfeldBC);
	end

	JTv = zeros(Float64,n)
	Vdatashape = reshape(v,nrec,nsrc);

	for k_batch = 1:numBatches
		batchIdxs = (k_batch-1)*batchSize + 1 : min(k_batch*batchSize,nsrc);
		V = P*Vdatashape[:,batchIdxs];
		V,Ainv = solveLinearSystem(H,V,Ainv,1);
		if pFor.useFilesForFields
			filename = getFieldsFileName(omega);
			file     = matopen(filename);
			U        = read(file,string("Ubatch_",k_batch));
			close(file);
		else
			U = pFor.Fields[:,batchIdxs];
		end
		# V = conj((1.0.+1im*vec(gamma)).*U).*V;
		V = conj((1.0.-1im*vec(gamma./omega)).*U).*V;

		JTv .+= omega^2*vec(real(sum(V,dims=2)));
	end

	if isa(Ainv,ShiftedLaplacianMultigridSolver)
		# println("Clearing!!!");
		clear!(Ainv.MG);
	end

	# V = P*Vdatashape;
	# V,Ainv = solveLinearSystem(H,V,Ainv,1)   # Lam = ForwardSolver\(P*V);

	# # JTv    = conj((1-1im*vec(gamma./omega)).*U).*V;
	# # V      = 0;
	# # JTv	 = omega^2*vec(sum(real(JTv),2));

    return JTv
end


using NNlib: conv, maxpool, meanpool, relu, softmax
using ChainRulesCore

# Initialization utils

xavier_uniform(T::Type, dims...; fan_in, fan_out) = sqrt(T(6) / (fan_in + fan_out)) * (2 * rand(T, dims...) .- 1)
xavier_uniform(dims...; fan_in, fan_out) = xavier_uniform(Float32, dims..., fan_in=fan_in, fan_out=fan_out)

# Linear layers

struct Linear{TW, Tb} <: Model
    W::TW
    b::Tb
end

function Linear(T::Type, dims::Pair)
    W = xavier_uniform(T, dims[2], dims[1], fan_in=dims[1], fan_out=dims[2])
    b = zeros(T, dims[2])
    return Linear(W, b)
end

Linear(dims::Pair) = Linear(Float32, dims)

function (l::Linear)(x)
    size_x = size(x)
    x_reshaped = reshape(x, size_x[1], :)
    out_reshaped = l.W * x_reshaped .+ l.b
    out = reshape(out_reshaped, :, size_x[2:end]...)
    return out
end

# Convolutional layers

struct Conv{Tw, Tb} <: Model
    w::Tw
    b::Tb
end

function Conv(T::Type, num_channels::Pair, filter_size::Tuple)
    fan_in = prod(filter_size) * num_channels[1]
    fan_out = prod(filter_size) * num_channels[2]
    w = xavier_uniform(T, filter_size..., num_channels..., fan_in=fan_in, fan_out=fan_out)
    b = zeros(T, 1, 1, num_channels[2])
    return Conv(w, b)
end

Conv(num_channels::Pair, filter_size::Tuple) = Conv(Float32, num_channels, filter_size)

(c::Conv)(x) = conv(x, c.w) .+ c.b

# Composition of layers

struct Compose{T <: Tuple} <: Model
    layers::T
end

Compose(first, rest...) = Compose((first, rest...))

function (m::Compose)(x)
    y = x
    for layer in m.layers
        y = layer(y)
    end
    return y
end

# Parameter-free layers (activations, pooling, reshaping...)

struct ReLU <: Model end

(::ReLU)(x) = relu.(x)

struct SoftMax{D} <: Model end

SoftMax(d=1) = SoftMax{d}()

(::SoftMax{D})(x) where D = softmax(x, dims=D)

struct MaxPool{S} <: Model end

MaxPool(s) = MaxPool{s}()

(p::MaxPool{S})(x) where S = maxpool(x, S)

struct MeanPool{S} <: Model end

MeanPool(s) = MeanPool{s}()

(p::MeanPool{S})(x) where S = meanpool(x, S)

struct Reshape{S} <: Model end

Reshape(s) = Reshape{s}()

(r::Reshape{S})(x) where S = reshape(x, S)

# Dropout

_dropout_shape(s, ::Colon) = size(s)
_dropout_shape(s, dims) = tuple((i in dims ? si : 1 for (i, si) in enumerate(size(s)))...)

function _dropout_mask(x::AbstractArray{T}, p; dims=:) where T
    y = rand(eltype(x), _dropout_shape(x, dims)...)
    y .= ifelse.(y .> p, T(1/(1 - p)), T(0))
    return y
end

function dropout(x, p; dims=:)
    if !is_training_mode()
        return x
    end
    y = _dropout_mask(x, p, dims=dims)
    return x .* y
end

function ChainRulesCore.rrule(::typeof(dropout), x, p; dims=:)
    if !is_training_mode()
        return x, c -> (ChainRulesCore.NoTangent(), c, ChainRulesCore.ZeroTangent())
    end
    y = _dropout_mask(x, p, dims=dims)
    return x .* y, c -> (ChainRulesCore.NoTangent(), c .* y, ChainRulesCore.ZeroTangent())
end

struct Dropout{P, D} <: Model end

Dropout(p, dims=:) = Dropout{p, dims}()

(d::Dropout{P, D})(x) where {P, D} = dropout(x, P, dims=D)


type TwoFactorShortRateTree{S <: ShortRateDynamics} <: ShortRateTree
  tree1::TrinomialTree
  tree2::TrinomialTree
  dynamics::S
  treeLattice::TreeLattice2D

  function TwoFactorShortRateTree{S}(tree1::TrinomialTree, tree2::TrinomialTree, dyn::S)
    twoFactorTree = new(tree1, tree2, dyn)
    twoFactorTree.treeLattice = TreeLattice2D(tree1, tree2, dyn.correlation, twoFactorTree)

    return twoFactorTree
  end
end

function tree(model::TwoFactorModel, grid::TimeGrid)
  #stuff
  dyn = get_dynamics(model)

  tree1 = TrinomialTree(dyn.xProcess, grid)
  tree2 = TrinomialTree(dyn.yProcess, grid)

  return TwoFactorShortRateTree{typeof(dyn)}(tree1, tree2, dyn)
end

get_size(tr::TwoFactorShortRateTree, i::Int) = get_size(tr.treeLattice, i)

descendant(tr::TwoFactorShortRateTree, i::Int, idx::Int, branch::Int) = descendant(tr.treeLattice, i, idx, branch)
probability(tr::TwoFactorShortRateTree, i::Int, idx::Int, branch::Int) = probability(tr.treeLattice, i, idx, branch)

function discount(tr::TwoFactorShortRateTree, i::Int, idx::Int)
  modulo = get_size(tr.tree1, i)
  new_idx = idx - 1

  index1 = round(Int, floor(new_idx % modulo)) + 1
  index2 = round(Int, floor(new_idx / modulo)) + 1

  x = get_underlying(tr.tree1, i, index1)
  y = get_underlying(tr.tree2, i, index2)

  r = short_rate(tr.dynamics, tr.treeLattice.tg.times[i], x, y)

  # println("$i $idx $r $x $y $modulo $index1 $index2")

  return exp(-r * tr.treeLattice.tg.dt[i])
end

type G2Dynamics{P <: Parameter} <: ShortRateDynamics
  fitting::P
  xProcess::OrnsteinUhlenbeckProcess
  yProcess::OrnsteinUhlenbeckProcess
  correlation::Float64
end

G2Dynamics{P <: Parameter}(fitting::P, a::Float64, sigma::Float64, b::Float64, eta::Float64, rho::Float64) =
          G2Dynamics{P}(fitting, OrnsteinUhlenbeckProcess(a, sigma), OrnsteinUhlenbeckProcess(b, eta), rho)

short_rate(dyn::G2Dynamics, t::Float64, x::Float64, y::Float64) = dyn.fitting(t) + x + y

type G2{AffineModelType, T <: TermStructure} <: TwoFactorModel{AffineModelType}
  modT::AffineModelType
  a::ConstantParameter
  sigma::ConstantParameter
  b::ConstantParameter
  eta::ConstantParameter
  rho::ConstantParameter
  phi::G2FittingParameter
  ts::T
  privateConstraint::PrivateConstraint
  common::ModelCommon
end

function G2{T <: TermStructure}(ts::T, a::Float64 = 0.1, sigma::Float64 = 0.01, b::Float64 = 0.1, eta::Float64 = 0.01, rho::Float64 = -0.75)
  a_const = ConstantParameter([a], PositiveConstraint())
  sigma_const = ConstantParameter([sigma], PositiveConstraint())
  b_const = ConstantParameter([b], PositiveConstraint())
  eta_const = ConstantParameter([eta], PositiveConstraint())
  rho_const = ConstantParameter([rho], BoundaryConstraint(-1.0, 1.0))

  privateConstraint = PrivateConstraint(ConstantParameter[a_const, sigma_const, b_const, eta_const, rho_const])

  phi = G2FittingParameter(a, sigma, b, eta, rho, ts)

  return G2{AffineModelType, T}(AffineModelType(), a_const, sigma_const, b_const, eta_const, rho_const, phi, ts, privateConstraint, ModelCommon())
end

get_eta(m::G2) = m.eta.data[1]
get_rho(m::G2) = m.rho.data[1]

get_params(m::G2) = Float64[get_a(m), get_sigma(m), get_b(m), get_eta(m), get_rho(m)]

generate_arguments!(m::G2) = m.phi = G2FittingParameter(get_a(m), get_sigma(m), get_b(m), get_eta(m), get_rho(m), m.ts)

get_dynamics(m::G2) = G2Dynamics(m.phi, get_a(m), get_sigma(m), get_b(m), get_eta(m), get_rho(m))

function V(m::G2, t::Float64)
  expat = exp(-get_a(m) * t)
  expbt = exp(-get_b(m) * t)
  cx = get_sigma(m) / get_a(m)
  cy = get_eta(m) / get_b(m)
  valuex = cx * cx * (t + (2.0 * expat - 0.5 * expat * expat - 1.5) / get_a(m))
  valuey = cy * cy * (t + (2.0 * expbt - 0.5 * expbt * expbt - 1.5) / get_b(m))
  value = 2.0 * get_rho(m) * cx * cy * (t + (expat - 1.0) / get_a(m) + (expbt - 1.0) / get_b(m) - (expat * expbt - 1.0) / (get_a(m) + get_b(m)))

  return valuex + valuey + value
end

A(m::G2, t::Float64, T::Float64) = discount(m.ts, T) / discount(m.ts, t) * exp(0.5 * (V(m, T - t) - V(m, T) + V(m, t)))
B(m::G2, x::Float64, t::Float64) = (1.0 - exp(-x * t)) / x

type G2SwaptionPricingFunction
  a::Float64
  sigma::Float64
  b::Float64
  eta::Float64
  rho::Float64
  w::Float64
  startTime::Float64
  fixedRate::Float64
  sigmax::Float64
  sigmay::Float64
  rhoxy::Float64
  mux::Float64
  muy::Float64
  payTimes::Vector{Float64}
  A::Vector{Float64}
  Ba::Vector{Float64}
  Bb::Vector{Float64}
end

function G2SwaptionPricingFunction(model::G2, w::Float64, startTime::Float64, payTimes::Vector{Float64}, fixedRate::Float64)
  a = get_a(model)
  sigma = get_sigma(model)
  b = get_b(model)
  eta = get_eta(model)
  rho = get_rho(model)

  sigmax = sigma * sqrt(0.5 * (-expm1(-2.0 * a * startTime)) / a)
  sigmay = eta * sqrt(0.5 * (-expm1(-2.0 * b * startTime)) / b)
  rhoxy = rho * eta * sigma * (-expm1(-(a + b) * startTime)) / ((a + b) * sigmax * sigmay)

  temp = sigma * sigma / (a * a)
  mux = -((temp + rho * sigma * eta / (a * b)) * (-expm1(-a * startTime)) - 0.5 * temp * (-expm1(-2.0 * a * startTime)) -
        rho * sigma * eta / (b * (a + b)) * (-expm1(-(b + a) * startTime)))

  temp = eta * eta / (b * b)
  muy = -((temp + rho * sigma * eta / (a * b)) * (-expm1(-b * startTime)) - 0.5 * temp * (-expm1(-2.0 * b * startTime)) -
        rho * sigma * eta / (a * (a + b)) * (-expm1(-(b + a) * startTime)))

  n = length(payTimes)
  A_ = zeros(n)
  Ba = zeros(n)
  Bb = zeros(n)
  for i = 1:n
    A_[i] = A(model, startTime, payTimes[i])
    Ba[i] = B(model, a, payTimes[i] - startTime)
    Bb[i] = B(model, b, payTimes[i] - startTime)
  end

  return G2SwaptionPricingFunction(a, sigma, b, eta, rho, w, startTime, fixedRate, sigmax, sigmay, rhoxy, mux, muy, payTimes, A_, Ba, Bb)
end

function operator(pricingFunc::G2SwaptionPricingFunction)
  phi = Normal()
  n = length(pricingFunc.payTimes)
  function _inner(x::Float64)
    temp = (x - pricingFunc.mux) / pricingFunc.sigmax
    txy = sqrt(1.0 - pricingFunc.rhoxy * pricingFunc.rhoxy)
    lambda = zeros(n)

    for i = 1:n
      tau = i == 1 ? pricingFunc.payTimes[1] - pricingFunc.startTime : pricingFunc.payTimes[i] - pricingFunc.payTimes[i - 1]
      c = i == n ? (1.0 + pricingFunc.fixedRate * tau) : pricingFunc.fixedRate * tau
      lambda[i] = c * pricingFunc.A[i] * exp(-pricingFunc.Ba[i] * x)
    end

    func = SolvingFunction(lambda, pricingFunc.Bb)
    s1d = BrentSolver(1000)

    yb = solve(s1d, operator(func), 1e-6, 0.00, -100.0, 100.0)

    h1 = (yb - pricingFunc.muy) / (pricingFunc.sigmay * txy) - pricingFunc.rhoxy * (x - pricingFunc.mux) / (pricingFunc.sigmax * txy)
    val = cdf(phi, -pricingFunc.w * h1)

    for i = 1:n
      h2 = h1 + pricingFunc.Bb[i] * pricingFunc.sigmay * sqrt(1.0 - pricingFunc.rhoxy * pricingFunc.rhoxy)
      kappa = -pricingFunc.Bb[i] * (pricingFunc.muy - 0.5 * txy * txy * pricingFunc.sigmay * pricingFunc.sigmay * pricingFunc.Bb[i] +
              pricingFunc.rhoxy * pricingFunc.sigmay * (x - pricingFunc.mux) / pricingFunc.sigmax)

      val -= lambda[i] * exp(kappa) * cdf(phi, -pricingFunc.w * h2)
    end

    blah = exp(-0.5 * temp * temp) * val / (pricingFunc.sigmax * sqrt(2.0 * pi))
    return exp(-0.5 * temp * temp) * val / (pricingFunc.sigmax * sqrt(2.0 * pi))
  end

  return _inner
end

function gen_swaption(model::G2, swaption::Swaption, fixedRate::Float64, range::Float64, intervals::Int)
  settlement = reference_date(model.ts)
  dc = model.ts.dc
  startTime = year_fraction(dc, settlement, swaption.swap.args.floatingResetDates[1])
  w = swaption.swap.payer[2]

  fixedPayTimes = zeros(length(swaption.swap.args.fixedPayDates))
  for i = 1:length(fixedPayTimes)
    fixedPayTimes[i] = year_fraction(dc, settlement, swaption.swap.args.fixedPayDates[i])
  end
  func = G2SwaptionPricingFunction(model, w, startTime, fixedPayTimes, fixedRate)

  upper = func.mux + range * func.sigmax
  lower = func.mux - range * func.sigmax

  integrator = SegmentIntegral(intervals)
  return swaption.swap.nominal * w * discount(model.ts, startTime) * QuantLib.Math.operator(integrator, QuantLib.operator(func), lower, upper)
end

discount_bond(model::G2, t::Float64, T::Float64, x::Float64, y::Float64) = A(model, t, T) * exp(-B(model, get_a(model), (T - t)) * x - B(model, get_b(model), (T - t)) * y)
discount_bond(model::G2, now::Float64, maturity::Float64, factors::Vector{Float64}) = discount_bond(model, now, maturity, factors[1], factors[2])