using GenericLinearAlgebra
using LinearAlgebra
using SpecialFunctions
using Dates
using PolynomialRoots
using MAT
using Printf
using FFTW

function get_QSP_proj(eps,zeta_A,zeta_B,p_A,p_B,stop_eps,R_init = 64, R_max = 2^10, even_root = false, only_a = false,
                      if_polish = false, GSLW_out = true)
    #------------------------------------------------------------------------------------------------------------
    #
    # Input:
    #     eps: rounding precision, usually be set as eps=stop_eps/30
    #     zeta_A, zeta_B: coefficients of Laurent polynomial A and B, only need to provide coefficients with
    #                     non-negative degree
    #     p_A, p_B: parity of A and B, 0 -- even, 1 -- odd
    #     stop_eps: algorithm will stop if it find a control achieve error less than stop_eps on N-th
    #               roots of unity, where N is degree of A
    #     R_init: number of bits used at the beginning
    #     R_max: number of bits available
    #     even_root: whether 1-A^2-B^2 is an even polynomial
    #     only_a: whether B=0
    #     if_polish: whether polish the roots, it cost more time but may improve performance
    #
    # Output:
    #    If GSLW_out=0, the algorithm will outout a list of 2*2 matrices {P_j} such that <+|E_0...E_{2n}|+>
    #                   approximate A(t^2)+iB(t^2) for all t with |t|=1, where E_j(t)=tP_j+(I-P_j)/t
    #    else,          the algorithm will outout a list of factors {phi_j} such that <+|E_0...E_{n}|+>
    #                   approximate A(t)+iB(t) for all t with |t|=1, where P_j=e^{iZ\phi_j/2}|+><+|e^{-iZ\phi_j/2},
    #                   E_j(t)=tP_j+(I-P_j)/t (Here A and B should be converted from real polynomials P and Q
    #                   satisfying constraints given in GSLW)
    #    Besides, the algorithm will return the L^∞ error of such approximation on N-th roots of unity
    #
    #------------------------------------------------------------------------------------------------------------
    #
    # Reference:
    #     Jeongwan Haah
    #     Product Decomposition of Periodic Functions in Quantum Signal Processing
    #
    # Author: X.M
    # Version 1.0 .... 02/2020
    #
    #------------------------------------------------------------------------------------------------------------

    # Step 1: Find rounded polynomials a(z) and b(z) from (1-10eps)A(z) and (1-10eps)B(z)

    R = R_init
    while(true)
        if(R>=R_max)
            return Inf,[]
        end
        setprecision(BigFloat, R)
        eps = big.(eps)
        zeta_A = big.(zeta_A)
        zeta_B = big.(zeta_B)
        zeta_A_eps = (big.(1)-big.(10)*eps)*zeta_A
        zeta_B_eps = (big.(1)-big.(10)*eps)*zeta_B
        degree_N = length(zeta_A) # the true degree of A is in fact degree_N-1
        rational_round = eps/big.((degree_N-1))

        zeta_a = zeros(Complex{BigFloat},2*degree_N-1,1)
        zeta_b = zeros(Complex{BigFloat},2*degree_N-1,1)

        if(norm(zeta_A_eps[1])<rational_round)
            zeta_a[degree_N] = 0
        else
            zeta_a[degree_N] = zeta_A_eps[1]
        end
        if(norm(zeta_B_eps[1])<rational_round)
            zeta_b[degree_N] = 0
        else
            zeta_b[degree_N] = zeta_B_eps[1]
        end
        degree_n = 0
        for i=1:degree_N-1
            if(norm(zeta_A_eps[i+1])<rational_round)
                zeta_a[degree_N+i] = big.(0)
            else
                zeta_a[degree_N+i] = zeta_A_eps[i+1]
                degree_n = i
            end
            if(norm(zeta_B_eps[i+1])<rational_round)
                zeta_b[degree_N+i] = big.(0)
            else
                zeta_b[degree_N+i] = zeta_B_eps[i+1]
                degree_n = i
            end
            if(p_A==1)
                zeta_a[degree_N-i] = -zeta_a[degree_N+i]
            else
                zeta_a[degree_N-i] = zeta_a[degree_N+i]
            end
            if(p_B==1)
                zeta_b[degree_N-i] = -zeta_b[degree_N+i]
            else
                zeta_b[degree_N-i] = zeta_b[degree_N+i]
            end
        end

        # calculate the coefficitents of 1-a(z)^2-b(z)^2

        zeta_1ab = zeros(Complex{BigFloat},4*degree_N-3,1)
        degree_1ab = 0

        for i=0:2*degree_N-2
            coff = big.(0)
            for j=i+1:2*degree_N-1
                coff -= zeta_a[j]*zeta_a[2*degree_N+i-j]
                coff -= zeta_b[j]*zeta_b[2*degree_N+i-j]
            end
            if(norm(coff)!=0)
                degree_1ab = i
            end
            zeta_1ab[2*degree_N-1+i] = real(coff)
            zeta_1ab[2*degree_N-1-i] = real(coff)
        end
        zeta_1ab[2*degree_N-1] += 1
        zeta_1ab = copy(zeta_1ab[2*degree_N-1-degree_1ab:2*degree_N-1+degree_1ab])

    # Step2: Find all roots of 1-a(z)^2-b(z)^2 with accuarcy R,
    #        we use Julia's internal rountine to find all roots

        # when 1-a(z)^2-b(z)^2 has special structure, we may reduce provlem scale

        if(only_a)
            zeta_1a1 = copy(zeta_a[degree_N-degree_n:degree_N+degree_n])
            zeta_1a1[degree_n+1] -= 1.0
            zeta_1a2 = copy(zeta_a[degree_N-degree_n:degree_N+degree_n])
            zeta_1a2[degree_n+1] += 1.0
            if(even_root)
                zeta_1a1 = copy(zeta_1a1[1:2:end])
                zeta_1a2 = copy(zeta_1a2[1:2:end])
            end
            all_root1 = roots(-zeta_1a1, polish = if_polish, epsilon = big.(0.0))
            all_root2 = roots(zeta_1a2, polish = if_polish, epsilon = big.(0.0))

            if(even_root)
                all_root = zeros(Complex{BigFloat}, 2*(length(all_root1)+length(all_root2)))
                for i=1:length(all_root1)
                    tmpnorm = norm(all_root1[i])
                    tmpangle = angle(all_root1[i])
                    all_root[2*i-1] = sqrt(tmpnorm)*exp(1im*tmpangle/2)
                    all_root[2*i] = -sqrt(tmpnorm)*exp(1im*tmpangle/2)
                end
                tmp_len = 2*length(all_root1)
                for i=1:length(all_root2)
                    tmpnorm = norm(all_root2[i])
                    tmpangle = angle(all_root2[i])
                    all_root[tmp_len+2*i-1] = sqrt(tmpnorm)*exp(1im*tmpangle/2)
                    all_root[tmp_len+2*i] = -sqrt(tmpnorm)*exp(1im*tmpangle/2)
                end
            else
                all_root = [all_root1;all_root2]
            end
        elseif(even_root)
            zeta_1ab2 = copy(zeta_1ab[1:2:end])
            all_rootmp = roots(zeta_1ab2, polish = if_polish, epsilon = big.(0.0))
            all_root = zeros(Complex{BigFloat}, 2*length(all_rootmp))
            for i=1:length(all_rootmp)
                tmpnorm = norm(all_rootmp[i])
                tmpangle = angle(all_rootmp[i])
                all_root[2*i-1] = sqrt(tmpnorm)*exp(1im*tmpangle/2)
                all_root[2*i] = -sqrt(tmpnorm)*exp(1im*tmpangle/2)
            end
        else
            all_root = roots(zeta_1ab, polish = if_polish, epsilon = big.(0.0))
        end

    # Step3: Evaluate complementary polynomials c(z) and d(z) on points {exp(2*k*pi*im/D)|k=1,2,...},
    #        where D is a power of 2 such that D>2n+1. These values will be utilized to perform FFT.

        # find roots with norm less than 1

        if(if_polish)
            eps_root = 1e-16
            root_list = zeros(Complex{BigFloat},4*degree_1ab,1)
            count = 0
            for i=1:length(all_root)
                norm_root = norm(all_root[i])
                if(abs(norm_root-1)<eps/(4*degree_N^2))
                    @printf("Warning: a root has norm close to 1\n")
                end
                if(norm_root<big.(1))
                    if(count<=degree_1ab)
                        flag = 0
                        for j=1:count
                            if(norm(root_list[j]-all_root[i])<eps_root)
                                flag = 1
                                break
                            end
                        end
                        if(flag==1)
                            continue
                        end

                        if(abs(all_root[i])<eps_root)
                            count += 1
                            root_list[count] = all_root[i]
                            continue
                        end
                        if(abs(imag(all_root[i]))<eps_root)
                            count += 2
                            root_list[count-1] = all_root[i]
                            root_list[count] = -all_root[i]
                            continue
                        end
                        if(abs(real(all_root[i]))<eps_root)
                            count += 2
                            root_list[count-1] = all_root[i]
                            root_list[count] = conj(all_root[i])
                            continue
                        end
                        count += 4
                        root_list[count-3] = all_root[i]
                        root_list[count-2] = conj(all_root[i])
                        root_list[count-1] = -all_root[i]
                        root_list[count] = conj(-all_root[i])
                    else
                        break
                    end
                end
            end
        else
            root_list = zeros(Complex{BigFloat},4*degree_1ab,1)
            count = 0
            for i=1:length(all_root)
                norm_root = norm(all_root[i])
                if(abs(norm_root-1)<eps/(4*degree_N^2))
                    @printf("Warning: a root has norm close to 1\n")
                end
                if(norm_root<big.(1))
                    count += 1
                    root_list[count] = all_root[i]
                end
            end
        end
        root_list = root_list[1:degree_1ab]
        count = count+1

        # find alpha, which is required in order to construct c(z) and d(z)

        z = big.(1)
        e_1 = big.(1)
        e_2 = big.(1)
        for j=1:degree_1ab
            e_1 *= (z-root_list[j])
            e_2 *= (big.(1)/z-root_list[j])
        end

        alpha = sum(zeta_1ab)/(e_1*e_2)
        alpha = real(alpha)

        if((alpha==NaN)||(alpha==Inf)||(alpha<0))
            @printf("Alpha is not correct, increase R\n")
            R = R*2
            continue
        end

        if(count<=degree_1ab)
            @printf("Number of roots is incorrect, increase R\n")
            R = R*2
            continue
        end

        # calculate value on D points

        D = 2^(ceil(Int,log2(2*degree_n+1)))
        cvalue_list = big.(zeros(Complex{BigFloat},D,1))
        dvalue_list = big.(zeros(Complex{BigFloat},D,1))
        for i=0:D-1
            z = exp(big.(2)*big.(pi)*im*big.(i)/big.(D))
            z_ = big.(1)/z

            e_1 = z_^floor(Int,degree_1ab/2)
            e_2 = z^floor(Int,degree_1ab/2)
            for j=1:degree_1ab
                e_1 *= (z-root_list[j])
                e_2 *= (z_-root_list[j])
            end
            dvalue_list[i+1] = sqrt(alpha)*(e_1+e_2)/big.(2)
            cvalue_list[i+1] = sqrt(alpha)*(e_1-e_2)/big.(-2)*big.(1)im
        end

    # Step 4: Compute the discrete fast Fourier transform of the function F(z) = a(z)I+b(z)iX+
    #         c(z)iY+d(z)iZ, where X,Y,Z are pauli matrices

        # calculate value on D points

        avalue_list = big.(zeros(Complex{BigFloat},D,1))
        bvalue_list = big.(zeros(Complex{BigFloat},D,1))
        for i=0:D-1
            z = exp(big.(2)*big.(pi)*im*big.(i)/big.(D))
            a_tmp = big.(0)
            b_tmp = big.(0)
            for j=1:2*degree_N-1
                a_tmp += z^(j-degree_N)*zeta_a[j]
                b_tmp += z^(j-degree_N)*zeta_b[j]
            end
            avalue_list[i+1] = a_tmp
            bvalue_list[i+1] = b_tmp
        end

        # perform FFT on each component

        C = zeros(Complex{BigFloat},2,2,2*degree_n+1)
        C_11 = BigFloatFFT(avalue_list+im*dvalue_list)
        C_12 = BigFloatFFT(cvalue_list+im*bvalue_list)
        C_21 = BigFloatFFT(-cvalue_list+im*bvalue_list)
        C_22 = BigFloatFFT(avalue_list-im*dvalue_list)
        C[1,1,:] = [C_11[D-degree_n+1:D];C_11[1:degree_n+1]]/big.(D)
        C[1,2,:] = [C_12[D-degree_n+1:D];C_12[1:degree_n+1]]/big.(D)
        C[2,1,:] = [C_21[D-degree_n+1:D];C_21[1:degree_n+1]]/big.(D)
        C[2,2,:] = [C_22[D-degree_n+1:D];C_22[1:degree_n+1]]/big.(D)

    # Step 5: From matrix C, caluculate the projection matrix P (and thus our objective E)

        if(GSLW_out)
            C = C[:,:,1:2:end]
            phi = zeros(Float64,degree_n+1)
            P = zeros(Complex{BigFloat},2,2,degree_n+1)
            for i=1:degree_n
                tmp_P = C[:,:,end]'*C[:,:,end]
                tmp_Q = C[:,:,1]'*C[:,:,1]
                P_i = tmp_P/tr(tmp_P)
                Q_i = tmp_Q/tr(tmp_Q)
                P[:,:,degree_n+2-i] = P_i
                phi[degree_n+2-i] = getarg(P[1,2,degree_n+2-i])
                C_new = zeros(Complex{BigFloat},2,2,degree_n+1-i)
                for j=1:degree_n+1-i
                    C_new[:,:,j] = C[:,:,j]*Q_i+C[:,:,j+1]*P_i
                end
                C = copy(C_new)
            end
            P[:,:,1] = copy(C[:,:,1])
            phi[1] = getarg(P[1,1,1])

            # Step 6: Check if phi is accurate enough, if so output phi, else increase R

            max_err = 0
            zeta_A = big.(zeta_A)
            zeta_B = big.(zeta_B)
            for i=1:degree_N
                t = exp(im*big.(pi)*big.(2*i)/big.(degree_N))
                QSP_value = getQSPvalue(P,t)
                true_value = big.(0)
                for j=0:length(zeta_A)-1
                    true_value += t^(j)*zeta_A[j+1]
                    true_value += im*t^(j)*zeta_B[j+1]
                    if(j!=0)
                        if(p_A==0)
                            true_value += t^(-j)*zeta_A[j+1]
                        else
                            true_value -= t^(-j)*zeta_A[j+1]
                        end
                        if(p_B==0)
                            true_value += im*t^(-j)*zeta_B[j+1]
                        else
                            true_value -= im*t^(-j)*zeta_B[j+1]
                        end
                    end
                end
                t_err = norm(QSP_value-true_value)
                if(t_err>max_err)
                    max_err = t_err
                end
            end
            @printf("For degree N = %d, precision R = %d, the estimated inf norm of err is %5.4e\n",degree_N-1,R,max_err)
            if(max_err<stop_eps)
                return max_err,phi
            else
                @printf("Error is too big, increase R.\n")
            end
            R = R*2
        else
            P = zeros(Complex{BigFloat},2,2,2*degree_n+1)
            for i=1:2*degree_n
                tmp_P = C[:,:,end]'*C[:,:,end]
                tmp_Q = C[:,:,1]'*C[:,:,1]
                P_i = tmp_P/tr(tmp_P)
                Q_i = tmp_Q/tr(tmp_Q)
                P[:,:,2*degree_n+2-i] = P_i
                C_new = zeros(Complex{BigFloat},2,2,2*degree_n+1-i)
                for j=1:2*degree_n+1-i
                    C_new[:,:,j] = C[:,:,j]*Q_i+C[:,:,j+1]*P_i
                end
                C = copy(C_new)
            end
            P[:,:,1] = copy(C[:,:,1])

            # Step 6: Check if P is accurate enough, if so output P, else increase R

            max_err = 0
            zeta_A = big.(zeta_A)
            zeta_B = big.(zeta_B)
            for i=1:degree_N
                t = exp(im*big.(pi)*big.(2*i)/big.(degree_N))
                QSP_value = getQSPvalue(P,t)
                true_value = big.(0)
                for j=0:length(zeta_A)-1
                    true_value += t^(2*j)*zeta_A[j+1]
                    true_value += im*t^(2*j)*zeta_B[j+1]
                    if(j!=0)
                        if(p_A==0)
                            true_value += t^(-2*j)*zeta_A[j+1]
                        else
                            true_value -= t^(-2*j)*zeta_A[j+1]
                        end
                        if(p_B==0)
                            true_value += im*t^(-2*j)*zeta_B[j+1]
                        else
                            true_value -= im*t^(-2*j)*zeta_B[j+1]
                        end
                    end
                end
                t_err = norm(QSP_value-true_value)
                if(t_err>max_err)
                    max_err = t_err
                end
            end
            @printf("For degree N = %d, precision R = %d, the estimated inf norm of err is %5.4e\n",degree_N-1,R,max_err)
            if(max_err<stop_eps)
                return max_err,P
            else
                @printf("Error is too big, increase R.\n")
            end
            R = R*2
        end
    end
end

function getarg(x)
    # Obtain the angle of a complex number
    x = x/abs(x)
    xr = acos(real(x))
    xi = asin(imag(x))
    if(xi>0)
        if(xr>0)
            return xr-2*pi
        else
            return xr
        end
    else
        return -xr
    end
end

function BigFloatFFT(x)
    # Perform FFT on vector x
    # This function only works for length(x) = 2^k
    N = length(x);
    xp = x[1:2:end];
    xpp = x[2:2:end];
    if(N>=8)
        Xp = BigFloatFFT(xp);
        Xpp = BigFloatFFT(xpp);
        X = zeros(Complex{BigFloat},N,1);
        Wn = exp.(big.(-2)im*big.(pi)*(big.(0:N/2-1))/big.(N));
        tmp = Wn .* Xpp;
        X = [(Xp + tmp);(Xp -tmp)];
    elseif(N==2)
        X = big.([1 1;1 -1])*x;
    elseif(N==4)
        X = big.([1 0 1 0; 0 1 0 -1im; 1 0 -1 0;0 1 0 1im]*[1 0 1 0;1 0 -1 0;0 1 0 1;0 1 0 -1])*x;
    end
    return X
end

function getQSPvalue(P,t)
    # Given P_j (0\le j \le n), yield <+|E_0...E_{n}|+>, where E_j(t)=tP_j+(I-P_j)/t
    eig_x = Complex{BigFloat}[1;1]/sqrt(big.(2))
    init = copy(eig_x)
    for i=1:size(P)[3]-1
        init = init/t+(t-big.(1)/t)*P[:,:,end+1-i]*init
    end
    return eig_x'*P[:,:,1]*init
end

# Test case 1: Hamiltonian simulation
#
# Here we want to approxiamte e^{-i\tau x} by Jacobi-Anger expansion:
#
# e^{-i\tau x} = J_0(\tau)+2\sum_{k even} (-1)^{k/2}J_{k}(\tau)T_k(x)+2i\sum_{k odd} (-1)^{(k-1)/2}J_{k}(\tau) T_k(x)
#
# We truncate the series up to N = 1.4\tau+log(10^{14}), which gives an polynomial approximation of e^{-i\tau x} with
# accuracy 10^{-14}. Besides, we deal with real and imaginary part of the truncated series seperatly and divide it
# by a constant factor 2 to enhance stability.
#
# A real polynomial P defined on [-1,1] is converted to input of the Haah method through the following way. Suppose
# P has expansion under Chebyshev basis, P = \sum_{k=0}^n a_kT_k(x), then A is taken as A=a_0+\sum_{k=1}^n a_k(z^k+1/z^k)/2
#
# parameters
#     stop_eps: desired accuracy
#     tau: the duration \tau in Hamiltonian simulation
#     R_init: number of bits used at the beginning
#     R_max: number of bits available

stop_eps = 1e-12
tau = 100
R_init = 1024
R_max = 1025

#------------------------------------------------------------------

phi1 = []
phi2 = []

for parity=0:1
    N = ceil.(Int,tau*1.4.+log(1e14))
    if(parity==0)
        setprecision(BigFloat,1024)
        if(mod(N,2)==1)
            N -= 1
        end
        zeta_A = zeros(Complex{BigFloat},1,N+1)
        zeta_B = zeros(Complex{BigFloat},1,N+1)
        p_A = 0
        p_B = 0
        for kk=1:(round(Int,N/2)+1)
            zeta_A[2*kk-1] = (-1)^(kk-1)*besselj(big.(2.0*(kk-1)),tau)/2.0
        end
    else
        setprecision(BigFloat,1024)
        if(mod(N,2)==0)
            N += 1
        end
        zeta_A = zeros(Complex{BigFloat},1,N+1)
        zeta_B = zeros(Complex{BigFloat},1,N+1)
        p_A = 0
        p_B = 0
        for kk=1:round(Int,(N+1)/2)
            zeta_A[2*kk] = (-1)^(kk-1)*besselj(big.(2.0*kk-1),tau)/2.0
        end
    end
    eps = stop_eps/30
    start_time = time()

    if(parity==0)
        err, phi1 = get_QSP_proj(eps,zeta_A,zeta_B,p_A,p_B,stop_eps,R_init,R_max,true,true)
    else
        err, phi2 = get_QSP_proj(eps,zeta_A,zeta_B,p_A,p_B,stop_eps,R_init,R_max,false,true)
    end

    elpased_time = time()-start_time
    @printf("Elapsed time is %4.2e s\n", elpased_time)
end

# Test case 2: Eigenstate filter
#
# Here we want to generate factors for the eigenstate filter function:
#
# f_n(x,\delta)=\frac{T_n(-1+2\frac{x^2-\delta^2}{1-\delta^2})}{T_n(-1+2\frac{-\delta^2}{1-\delta^2})}.
#
# We divide f_n by a constant factor \sqrt{2} to enhance stability.
#
# A real polynomial P defined on [-1,1] is converted to input of the Haah method through the following way. Suppose
# P has expansion under Chebyshev basis, P = \sum_{k=0}^n a_kT_k(x), then A is taken as A=a_0+\sum_{k=1}^n a_k(z^k+1/z^k)/2
#
# Reference: Lin Lin and Yu Tong
#            Solving quantum linear system problem with near-optimal complexity
#
# parameters
#     stop_eps: desired accuracy
#     n, \delta: parameters of f_n
#     R_init: number of bits used at the beginning
#     R_max: number of bits available
#

stop_eps = 1e-12
n = 100
delta = 0.03
R_init = 1024
R_max = 2048

#------------------------------------------------------------------

function f_n(x,n,delta)
    val = copy(x)
    delta = big.(delta)
    fact = chebyshev(-big.(1)-big.(2)*delta^2/(big.(1)-delta^2),n)
    if(length(x)==1)
        return chebyshev(-big.(1)+big.(2)*(x^2-delta^2)/(big.(1)-delta^2),n)/fact
    else
        for i=1:length(x)
            val[i] = chebyshev(-1+2*(x[i]^2-delta^2)/(1-delta^2),n)/fact
        end
        return val
    end
end

function chebyshev(x,n) # T_n(x)
    if(abs(x)<=1)
        return cos(big.(n)*acos(x))
    elseif(x>1)
        return cosh(big.(n)*acosh(x))
    else
        return big.((-1)^n)*cosh(big.(n)*acosh(-x))
    end
end

setprecision(BigFloat,1024)

# Obtain expansion of f_n under Chebyshev basis via FFT

M = 2*n
theta = range(0, stop=2*pi, length=2*M+1)
theta = theta[1:2*M]
f = f_n(-cos.(theta),n,delta)
c = real(fft(f))
c = c[1:M+1]
c[2:end-1] = c[2:end-1]*2
c[2:2:end] = - c[2:2:end]
c = c / (2*M)

zeta_A = zeros(Complex{BigFloat},1,length(c))
zeta_B = zeros(Complex{BigFloat},1,length(c))
zeta_A[1] = c[1]/sqrt(big.(2.0))
for k=2:length(c)
    if(mod(k,2)==1)
        zeta_A[k] = c[k]/sqrt(big.(8.0))
    end
end
p_A = 0
p_B = 0
eps = stop_eps/30
start_time = time()
err,phi = get_QSP_proj(eps,zeta_A,zeta_B,p_A,p_B,stop_eps,R_init,R_max,true,true)
elpased_time = time()-start_time
@printf("Elapsed time is %4.2e s\n", elpased_time)

# Test case 3: Matrix inversion
#
# We would like to approximate 1/x over [1/kappa,1] by a polynomial, such polynomial is generated
# by Remez algorithm
#
# A real polynomial P defined on [-1,1] is converted to input of the Haah method through the following way. Suppose
# P has expansion under Chebyshev basis, P = \sum_{k=0}^n a_kT_k(x), then A is taken as A=a_0+\sum_{k=1}^n a_k(z^k+1/z^k)/2
#
# parameters
#     stop_eps: desired accuracy
#     kappa: parameters of polynomial approximation
#     R_init: number of bits used at the beginning
#     R_max: number of bits available
#

stop_eps = 1e-12
kappa = 20
R_init = 2048
R_max = 2049

#------------------------------------------------------------------
# even approximation

eps = stop_eps/30

# enter your path here
matpath2 = "Data\\inversex\\"
vars = matread(matpath2 * "coef_xeven_" * string(kappa)*"_6"* ".mat")
coef = vars["coef"]
zeta_A = zeros(Complex{BigFloat},1,2*length(coef)-1)
zeta_B = zeros(Complex{BigFloat},1,2*length(coef)-1)
for kk=1:length(coef)
    zeta_A[2*kk-1] = coef[kk]/2
end
zeta_A[1] *= 2
p_A = 0
p_B = 0
start_time = time()
err,phi1 = get_QSP_proj(eps,zeta_A,zeta_B,p_A,p_B,stop_eps,R_init,R_max,true,true)
elpased_time = time()-start_time
@printf("Elapsed time is %4.2e s\n", elpased_time)

# odd approximation

vars = matread(matpath2 * "coef_xodd_" * string(kappa)*"_6"* ".mat")
coef = vars["coef"]
zeta_A = zeros(Complex{BigFloat},1,2*length(coef))
zeta_B = zeros(Complex{BigFloat},1,2*length(coef))
for kk=1:length(coef)
    zeta_A[2*kk] = coef[kk]/2
end
p_A = 0
p_B = 0
start_time = time()
err,phi2 = get_QSP_proj(eps,zeta_A,zeta_B,p_A,p_B,stop_eps,R_init,R_max,false,true)
elpased_time = time()-start_time
@printf("Elapsed time is %4.2e s\n", elpased_time)


using BNPRegressionGGA2021
using CSV
using DataFrames
using Random
const BNP = BNPRegressionGGA2021

# Simulate the data
begin
    include("dgp-normal-normal1.jl");
    Random.seed!(1);
    N0, N1 = 500, 50;
    dy, y0, x0, X0, y1, x1, X1 = simulate_sample(N0, N1);
    m0, o0 = mean(y0), std(y0)
    y0 = (y0 .- m0) ./ o0
end;
# did not work with 1000
# did not work with 2000
# did not work with 3000
# did not work with 4000
# worked better with 5000
# worked better with 6000
# worked with 10000

# Fit the model
begin
    update_γ = false;
    mapping = [[1], collect(2:size(X0, 2))];
    m = BNP.DGSBPNormal(; y0, X0, y1, X1, mapping, update_γ);
    chainf, chainβ, chaing = BNP.sample!(m; mcmcsize = 20000, burnin = 10000);
end;

# Save the results as a CSV file
begin
    df = DataFrame(
        y = y1,
        x = x1,
        fh = mean(chainf),
        f0 = pdf.(dy.(x1), m0 .+ o0 .* y1) .* o0
    )
    filename = "data/dgp-normal-normal-example1.csv"
    CSV.write(filename, df)
end


"""
    module ArnoldiMethodTransformations

Provides convenience wrapper for interfacing with the package ArnoldiMethod.

Implements the shift-and-invert transformation detailed [here](https://haampie.github.io/ArnoldiMethod.jl/stable/).

The main functions are `partialschur(A,[B],σ; kwargs...)` and `partialeigen(A,σ; kwargs...)`

The constants `USOLVER`, `PSOLVER`, and `MSOLVER` are exported.
"""
module ArnoldiMethodTransformations

export PSOLVER
export PSolver
export MSOLVER
export MSolver
export USOLVER
export USolver

using ArnoldiMethod
using LinearAlgebra
using LinearMaps
#using MUMPS
using Pardiso
using SparseArrays

abstract type AbstractSolver end
struct PSolver <: AbstractSolver end
struct MSolver <: AbstractSolver end
struct USolver <: AbstractSolver end

# solver constants (exported)
const PSOLVER = PSolver()
const MSOLVER = MSolver()
const USOLVER = USolver()

# colors for pretty printing
const DEFAULT_SOLVER = USOLVER
const SOLVER_COLOR = :cyan
const PRINTED_TYPE_COLOR = 171

# main structure used internally
struct ShiftAndInvert{M,TLU,TB,V,Σ,TZ}
    LU::TLU
    B::TB
    temp1::V
    temp2::V
    σ::Σ
    diag_inv_B::Bool
    issymmetric::Bool
    Z::TZ
end

# for ordinary EVP with Dense Matrix
function ShiftAndInvert(
            A::Matrix{TA};
            σ::Σ=0,
            lupack::AbstractSolver = DEFAULT_SOLVER
            ) where {TA,TB,Σ<:Number}

    diag_inv_B=true
    T = promote_type(TA,Σ)
    α = A-σ*I
    β = Matrix(one(T)I,size(A))

    temp1 = Vector{T}(undef,size(α,1))
    temp2 = Vector{T}(undef,size(α,1))
    issym = issymmetric(α) & issymmetric(β)
    Z = spzeros(T,size(β)...)

    M, a_lu = initialize_according_to_package(lupack,α,issym,T,temp1,temp2)
    return ShiftAndInvert{typeof(M),typeof(a_lu),typeof(β),typeof(temp1),typeof(σ),typeof(Z)}(a_lu,β,temp1,temp2,σ,diag_inv_B,issym,Z)
end
# for generalized EVP with Dense Matrices
function ShiftAndInvert(
            A::Matrix{TA},
            B::Matrix{TB};
            σ::Σ=0,
            diag_inv_B::Bool=false,
            lupack::AbstractSolver = DEFAULT_SOLVER
            ) where {TA,TB,Σ<:Number}

    T = promote_type(TA,TB,Σ)
    if diag_inv_B
        for i ∈ diagind(B) B[i]=1/B[i] end
        α = B*A-σ*I
        for i ∈ diagind(B) B[i]=1/B[i] end
        β = Matrix(one(T)I,size(B))
    else
        α = A-σ*B
        β = convert(Matrix{T},B)
    end
    temp1 = Vector{T}(undef,size(α,1))
    temp2 = Vector{T}(undef,size(α,1))
    issym = issymmetric(α) & issymmetric(β)
    Z = spzeros(T,size(β)...)

    M, a_lu = initialize_according_to_package(lupack,α,issym,T,temp1,temp2)
    return ShiftAndInvert{typeof(M),typeof(a_lu),typeof(β),typeof(temp1),typeof(σ),typeof(Z)}(a_lu,β,temp1,temp2,σ,diag_inv_B,issym,Z)
end
# for ordinary EVP with Sparse Matrix
function ShiftAndInvert(
            A::SparseMatrixCSC{TA};
            σ::Σ=0,
            lupack::AbstractSolver = DEFAULT_SOLVER
            ) where {TA,TB,Σ}

    diag_inv_B=true
    T = promote_type(TA,Σ)
    α = A-σ*I
    β = sparse(one(T)I,size(A))
    temp1 = Vector{T}(undef,size(α,1))
    temp2 = Vector{T}(undef,size(α,1))
    issym = issymmetric(α) & issymmetric(β)
    Z = spzeros(T,size(β)...)

    # initialize according to package used
    M, a_lu = initialize_according_to_package(lupack,α,issym,T,temp1,temp2)
    return ShiftAndInvert{typeof(M),typeof(a_lu),typeof(β),typeof(temp1),typeof(σ),typeof(Z)}(a_lu,β,temp1,temp2,σ,diag_inv_B,issym,Z)
end
# for generalized EVP with Sparse Matrices
function ShiftAndInvert(
            A::SparseMatrixCSC{TA},
            B::SparseMatrixCSC{TB};
            σ::Σ=0,
            diag_inv_B::Bool=false,
            lupack::AbstractSolver = DEFAULT_SOLVER
            ) where {TA,TB,Σ}

    T = promote_type(TA,TB,Σ)
    if diag_inv_B
        for i ∈ diagind(B) B[i]=1/B[i] end
        α = B*A-σ*I
        for i ∈ diagind(B) B[i]=1/B[i] end
        β = sparse(one(T)I,size(B))
    else
        α = A-σ*B
        β = convert(SparseMatrixCSC{T,Int},B)
    end
    temp1 = Vector{T}(undef,size(α,1))
    temp2 = Vector{T}(undef,size(α,1))
    issym = issymmetric(α) & issymmetric(β)
    Z = spzeros(T,size(β)...)

    # initialize according to package used
    M, a_lu = initialize_according_to_package(lupack,α,issym,T,temp1,temp2)
    return ShiftAndInvert{typeof(M),typeof(a_lu),typeof(β),typeof(temp1),typeof(σ),typeof(Z)}(a_lu,β,temp1,temp2,σ,diag_inv_B,issym,Z)
end
# if not both dense or both sparse
ShiftAndInvert(A,B;kwargs...) = throw("typeof(A)=$(typeof(A)), typeof(B)=$(typeof(B)). Either both sparse or both dense")

# define action of ShiftAndInvert
function (SI::ShiftAndInvert{M})(y,x) where M<:Union{USolver,MSolver}
    #first multiply Bv
    if SI.diag_inv_B
        for i ∈ eachindex(SI.temp1) SI.temp1[i] = x[i] end
    else
        mul!(SI.temp1, SI.B, x)
    end
    # then left-divide (A-σB)⁻¹*(Bv)
    ldiv!(y, SI.LU, SI.temp1)
    return nothing
end
function (SI::ShiftAndInvert{PSolver})(y,x)
    #first multiply Bv
    if SI.diag_inv_B
        for i ∈ eachindex(SI.temp1) SI.temp1[i] = x[i] end
    else
        mul!(SI.temp1, SI.B, x)
    end
    # then left-divide (A-σB)⁻¹*(Bv)
    pardiso(SI.LU,SI.temp2,SI.Z,SI.temp1)
    for i ∈ eachindex(y) y[i] = SI.temp2[i] end
    return nothing
end

# generate a solver-dependent LU object
# UMFPACK
initialize_according_to_package(::USolver,α,args...) = USOLVER, lu(α)
# MUMPS
function initialize_according_to_package(::MSolver,α,args...)
    throw("MUMPS still under construction")
    try MPI catch; throw(ErrorException("MUMPS not loaded. Try again after `using MUMPS`")) end
    MPI.Initialized() ? nothing : MPI.Init()
    return MSOLVER, mumps_factorize(α)
end
# PARDISO
function initialize_according_to_package(::PSolver,α,issym,type,x,y)
    try PardisoSolver catch; throw(ErrorException("Pardiso not loaded. Try again after `using Pardiso`")) end
    A_lu = PardisoSolver()
    set_iparm!(A_lu,1,1) # don't revert to defaults
    set_iparm!(A_lu,12,1) # transpose b/c of CSR vs SCS
    # x = Vector{type}(undef,size(α,1))
    # y = Vector{type}(undef,size(α,1))
    if issym & (type<:Real)
        set_matrixtype!(A_lu,Pardiso.REAL_SYM)
        pardiso(A_lu,x,sparse(tril(α)),y)
    elseif issym & (type<:Complex)
        set_matrixtype!(A_lu,Pardiso.COMPLEX_SYM)
        pardiso(A_lu,x,sparse(tril(α)),y)
    elseif !issym & (type<:Real)
        set_matrixtype!(A_lu,Pardiso.REAL_NONSYM)
        pardiso(A_lu,x,sparse(α),y)
    else # if !issym & type<:Complex
        set_matrixtype!(A_lu,Pardiso.COMPLEX_NONSYM)
        pardiso(A_lu,x,sparse(α),y)
    end
    set_phase!(A_lu,12) # analyze and factorize
    set_phase!(A_lu,33) # set to solve for future calls
    return PSOLVER, A_lu
end

# pretty printing
function Base.show(io::IO,si::ShiftAndInvert{M}) where M
    printstyled(io,"ShiftAndInvert ",color=PRINTED_TYPE_COLOR)
    print(io,"(using ")
    if M==MSolver
        printstyled(io,"MUMPS",color=SOLVER_COLOR)
    elseif M==PSolver
        printstyled(io,"Pardiso",color=SOLVER_COLOR)
    elseif M==USolver
        printstyled(io,"UMFPACK",color=SOLVER_COLOR)
    end
    print(io,")")
end

################################################################################

# add method to partialschur(::ShiftAndInvert)
function ArnoldiMethod.partialschur(si::ShiftAndInvert; kwargs...)
    a = LinearMap{eltype(si.B)}(si, size(si.B,1); ismutating=true, issymmetric=si.issymmetric)
    return partialschur(a; kwargs...)
end
# add method to partialschur(A,σ)
function ArnoldiMethod.partialschur(
            A::AbstractArray,
            σ::Number;
            lupack::AbstractSolver=DEFAULT_SOLVER,
            kwargs...)
    return partialschur(ShiftAndInvert(A; σ=σ, lupack=lupack); kwargs...)
end
# add method to partialschur(A,B,σ)
function ArnoldiMethod.partialschur(
            A::AbstractArray,
            B::AbstractArray,
            σ::Number;
            diag_inv_B::Bool = isdiag(B) && !any(iszero.(diag(B))),
            lupack::AbstractSolver=DEFAULT_SOLVER,
            kwargs...)
    return partialschur(ShiftAndInvert(A, B; σ=σ, diag_inv_B=diag_inv_B, lupack=lupack); kwargs...)
end

# add method to partialeigen(decomp,σ)
"""
    partialeigen(decomp, σ) -> λ, v

    Transforms a partial Schur decomposition into an eigendecomposition, outputting evals `λ` and evecs `v`.
    It undoes the shift-and-invert of the eigenvalues by `σ`.
"""
function ArnoldiMethod.partialeigen(
            decomp::ArnoldiMethod.PartialSchur,
            σ::Number)
    λ, v = partialeigen(decomp)
    foreach(i->λ[i]=σ+1/λ[i],eachindex(λ))
    return λ, v
end

################################################################################
"""
    struct ShiftAndInvert

Container for arrays in `ArnoldiMethod.partialschur`

-------------

    function ShiftAndInvert(A, [B]; σ=0, diag_inv_B=false, lupack=USOLVER) -> si

create a LinearMap object to feed to ArnoldiMethod.partialschur which transforms `Ax=λBx` into `(A-σB)⁻¹Bx=x/(λ-σ)`.

Set `diag_inv_B=true` if `B` is both diagonal and invertible, so that it is easy to compute `B⁻¹`. In this case instead return a linear map `(B⁻¹A-σI)⁻¹`, which has same evals as above.

`A` and `B` must both be sparse or both dense. `A`, `B`, `σ` need not have common element type.

Keyword `lupack` determines what linear algebra library to use. Options are `PSOLVER`, `MSOLVER`, `USOLVER` (the default).
For all solvers besides UMFPACK, the appropriate library must be loaded at the top level.
----------

    function ::ShiftAndInvert(y,x)

`A\\B*x = y`
"""
ShiftAndInvert

"""
    partialschur(A, [B], σ; [diag_inv_B=false, lupack=USOLVER, kwargs...]) -> decomp, history

Partial Schur decomposition of `A`, with shift `σ` and mass matrix `B`, solving `Av=σBv`

Keyword `diag_inv_B` defaults to `true` if `B` is both diagonal and invertible. This enables
a simplified shift-and-invert scheme.

Keyword `lupack` determines what linear algebra library to use. Options are `PSOLVER`, `MSOLVER`, `USOLVER` (the default).
For all solvers besides UMFPACK, the appropriate library must be loaded at the top level.

For other keywords, see ArnoldiMethod.partialschur

see also: [`partialeigen`](@ref) in ArnoldiMethod
"""
partialschur

end # module


# Evaluating S(1) as Manifold for Gaussian Products on manifold

using LinearAlgebra
using ApproxManifoldProducts

using TransformUtils

const TU = TransformUtils

# using Gadfly

logmap_SO2(Rl::Matrix{Float64}) = sign(Rl[2,1])*acos(Rl[1,1])
difftheta(wth1::Float64, wth2::Float64)::Float64 = logmap_SO2(TU.R(wth1)'*TU.R(wth2))

#assume an unwrapped circle

μ1 = # 0.0
μ2 = pi/2.0

Λ1 = 10.0
Λ2 = 10.0

Lambdas = [Λ1; Λ2]
mus = [μ1; μ2]

μ = get2DMuMin(mus, Lambdas, diffop=difftheta, initrange=(-pi+0.0,pi+0.0), periodicmanifold=wrapRad)[1]
TU.wrapRad(μ)

N = 1000

BUPS = Float64[get2DMuMin(mus, Lambdas, diffop=difftheta, periodicmanifold=wrapRad)[1] for i in 1:N]
# BUPS[:] .= TU.wrapRad.(BUPS)

@test 0.3*N < sum( 0 .< BUPS .< 3pi/2.0)
@test 0.3*N < sum( -3pi/2 .< BUPS .< 0.0)
@test sum( -0.1 .< BUPS .< 0.1) < 0.01*N

# plot(x=BUPS, Geom.histogram)

# @testset "fanning ternary means..." begin
#
# μ1 = 0.0
# μ2 = -4pi/6.0
# μ3 = 4pi/6.0
#
# Λ1 = 1.0
# Λ2 = 1.0
# Λ3 = 1.0
#
# Lambdas = [Λ1; Λ2; Λ3]
# mus = [μ1; μ2; μ3]
#
#
# μ = get2DMuMin(mus, Lambdas, diffop=difftheta, periodicmanifold=wrapRad, initrange=(-pi+0.0,pi+0.0))[1]
#
# # TU.wrapRad(μ)
#
# N = 100
# BUPS = Float64[get2DMuMin(mus, Lambdas, diffop=difftheta, periodicmanifold=wrapRad, initrange=(-pi+0.0,pi+0.0))[1] for i in 1:N]
# # BUPS[:] .= TU.wrapRad.(BUPS)
#
# # plot(x=BUPS, Geom.histogram)
#
# @test 0.20*N < sum( pi/2 .< BUPS .< 3pi/2.0)
# @test 0.20*N < sum( -3pi/2 .< BUPS .< -pi/2)
# @test 0.20*N < sum( -0.1 .< BUPS .< 0.1)
#
#
# end
#
# end
#
#
#
# @testset "basic linear 2D Rotation find mu and Lambdas Gaussian product test..." begin
#
# # expect \mu = 5pi/6
# μ1 = 4pi/6.0
# μ2 = -pi+0.0
# μ3 = -1.05
# # μ3 = -0.5
# #
# # μ1 = 0.0
# # μ2 = 2.0*pi/3.0
# # μ3 = -2.0*pi/3.0
# #
#
# Λ1 = 1.0
# Λ2 = 1.0
# Λ3 = 1.0
#
#
# mus = [μ1; μ2; μ3]
# Lambdas = [Λ1; Λ2; Λ3]
#
# difftheta(μ1, 0.0)
#
# difftheta(μ1, -pi+0.0)
# difftheta(μ1, pi+0.0)
#
# difftheta(μ1, -1.04)
# difftheta(μ1, -1.05)
#
#
# @test abs(difftheta(μ2, 0.0)) - pi < 1e-8
# @test abs(difftheta(μ2, 0.0001) ) - pi < 1e-3
# @test abs(difftheta(μ2, -0.0001)) - pi < 1e-3
#
# difftheta(μ2, -1.04)
# difftheta(μ2, -1.05)
#
#
# @test abs(difftheta(μ2, -pi+0.0)) < 1e-8
# @test abs(difftheta(μ2, pi+0.0) ) < 1e-8
#
# @test 1e-5 < abs( difftheta(μ1, 5pi/6.0) )
# @test 1e-5 < abs( difftheta(μ2, 5pi/6.0) )
# @test abs(difftheta(μ1, 5pi/6.0) + difftheta(μ2, 5pi/6.0)) <1e-8
#
# resid2DLinear(0.0, mus, Lambdas, diffop=difftheta)
#
#
#
# # @test abs(resid2DLinear(5pi/6.0, mus, Lambdas, diffop=difftheta)) < 1e-10
#
# resid2DLinear(μ1, mus, Lambdas, diffop=difftheta)
#
# resid2DLinear(pi+0.0, mus, Lambdas, diffop=difftheta)
# resid2DLinear(-pi, mus, Lambdas, diffop=difftheta)
#
# μ = get2DMu(mus, Lambdas, diffop=difftheta, periodicmanifold=wrapRad)
# Λ = get2DLambda(Lambdas)
#
# # @test abs(μ[1] - 8/5) < 1e-10
# #
# # @test abs(Λ - 5/4) < 1e-10
#
# end
#

#


using DualNumbers

struct DualQuaternion{T<:Real} <: Number
  q0::Quaternion{T}
  qe::Quaternion{T}
  norm::Bool
end

DualQuaternion{T}(dq::DualQuaternion) where {T<:Real} = DualQuaternion{T}(dq.q0, dq.qe, dq.norm)
function DualQuaternion{T}(d1::Dual, d2::Dual, d3::Dual, d4::Dual, n::Bool=false) where {T<:Real}
  return DualQuaternion{T}(
    Quaternion(DualNumbers.value(d1), DualNumbers.value(d2), DualNumbers.value(d3), DualNumbers.value(d4), n),
    Quaternion(DualNumbers.epsilon(d1), DualNumbers.epsilon(d2), DualNumbers.epsilon(d3), DualNumbers.epsilon(d4)),
    n,
  )
end
function DualQuaternion{T}(q0::Quaternion) where {T<:Real}
  return DualQuaternion{T}(convert(Quaternion{T}, q0), zero(Quaternion{T}), q0.norm)
end
function DualQuaternion{T}(d::Dual) where {T<:Real}
  return DualQuaternion(
    Quaternion{T}(DualNumbers.value(d)),
    Quaternion{T}(DualNumbers.epsilon(d)),
    (DualNumbers.value(d)==one(DualNumbers.value(d))) & iszero(DualNumbers.epsilon(d)))
end
function DualQuaternion{T}(x::Real) where {T<:Real}
  return DualQuaternion(convert(Quaternion{T}, x), zero(Quaternion{T}), abs(x) == one(x))
end

DualQuaternion(q0::Quaternion, qe::Quaternion) = DualQuaternion(promote(q0, qe)..., false)
DualQuaternion(d1::Dual, d2::Dual, d3::Dual, d4::Dual, n::Bool=false) =
  DualQuaternion(Quaternion(DualNumbers.value(d1), DualNumbers.value(d2), DualNumbers.value(d3), DualNumbers.value(d4), n),
                  Quaternion(DualNumbers.epsilon(d1), DualNumbers.epsilon(d2), DualNumbers.epsilon(d3), DualNumbers.epsilon(d4)), n)
DualQuaternion(x::Real) = DualQuaternion(Quaternion(x), Quaternion(zero(x)), abs(x) == one(x))
DualQuaternion(d::Dual) = DualQuaternion(Quaternion(DualNumbers.value(d)), Quaternion(DualNumbers.epsilon(d)), (DualNumbers.value(d)==one(DualNumbers.value(d))) & iszero(DualNumbers.epsilon(d)))
DualQuaternion(q::Quaternion) = DualQuaternion(q, zero(q), q.norm)
DualQuaternion(a::Vector) = DualQuaternion(zero(Quaternion{typeof(a[1])}), Quaternion(a))

const DualQuaternionF16 = DualQuaternion{Float16}
const DualQuaternionF32 = DualQuaternion{Float32}
const DualQuaternionF64 = DualQuaternion{Float64}

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

dualquat(q1, q2) = DualQuaternion(q1, q2)
dualquat(q1, q2, n) = DualQuaternion(q1, q2, n)
dualquat(d1, d2, d3, d4) = DualQuaternion(d1, d2, d3, d4)
dualquat(d1, d2, d3, d4, n) = DualQuaternion(d1, d2, d3, d4, n)
dualquat(x) = DualQuaternion(x)

function show(io::IO, dq::DualQuaternion)
  show(io, dq.q0)
  print(io, " + ( ")
  show(io, dq.qe)
  print(io, " )du")
end

Q0(dq::DualQuaternion) = dq.q0
Qe(dq::DualQuaternion) = dq.qe

(/)(dq::DualQuaternion, x::Real) = DualQuaternion(dq.q0 / x, dq.qe / x)

(/)(dq::DualQuaternion, d::Dual) =
  DualQuaternion(dual(dq.q0.s, dq.qe.s) / d,
                  dual(dq.q0.v1, dq.qe.v1) / d,
                  dual(dq.q0.v2, dq.qe.v2) / d,
                  dual(dq.q0.v3, dq.qe.v3) / d)

abs2(dq::DualQuaternion) = dq.norm ? dual(one(dq.q0.s)) :
  dual(abs2(dq.q0),
        2.0 * (dq.q0.s  * dq.qe.s  +
                dq.q0.v1 * dq.qe.v1 +
                dq.q0.v2 * dq.qe.v2 +
                dq.q0.v3 * dq.qe.v3))

abs(dq::DualQuaternion) = dq.norm ? dual(one(dq.q0.s)) : sqrt(abs2(dq))
float(dq::DualQuaternion{T}) where T = convert(DualQuaternion{float(T)}, dq)

conj(dq::DualQuaternion) = DualQuaternion(conj(dq.q0), conj(dq.qe), dq.norm)
dconj(dq::DualQuaternion) = DualQuaternion(dq.q0, -dq.qe, dq.norm)

inv(dq::DualQuaternion) = dq.norm ? conj(dq) : conj(dq) / abs2(dq)

function normalize(dq::DualQuaternion)
  if (dq.norm)
    return dq
  end
  a = abs(dq)
  if abs(a) > 0
    qa = dq / a
    dualquat(qa.q0, qa.qe, true)
  else
    dq
  end
end

function normalizea(dq::DualQuaternion{T}) where {T}
  if (dq.norm)
    return (dq, one(Dual{T}))
  end
  a = abs(dq)
  if abs(a) > 0
    qa = dq / a
    dualquat(qa.q0, qa.qe, true), a
  else
    dq, zero(Dual{T})
  end
end

(-)(dq::DualQuaternion) = DualQuaternion(-dq.q0, -dq.qe, dq.norm)

(+)(dq::DualQuaternion, dw::DualQuaternion) = DualQuaternion(dq.q0 + dw.q0, dq.qe + dw.qe)
(-)(dq::DualQuaternion, dw::DualQuaternion) = DualQuaternion(dq.q0 - dw.q0, dq.qe - dw.qe)
(*)(dq::DualQuaternion, dw::DualQuaternion) = DualQuaternion(dq.q0 * dw.q0,
                                                              dq.q0 * dw.qe + dq.qe * dw.q0,
                                                              dq.norm && dw.norm)
(*)(dq::DualQuaternion, d::Dual) = (*)(Base.promote(dq, d)...)
(*)(d::Dual, dq::DualQuaternion) = (*)(Base.promote(d, dq)...)
(/)(dq::DualQuaternion, dw::DualQuaternion) = dq * inv(dw)
(==)(q::DualQuaternion, w::DualQuaternion) = (q.q0 == w.q0) & (q.qe == w.qe) # ignore .norm field

function angleaxis(dq::DualQuaternion)
  tq = dq.qe * conj(dq.q0)
  t = [2.0 * tq.v1, 2.0 * tq.v2, 2.0 * tq.v3]
  q0s = dq.q0.s
  th0, s0 = angleaxis(dq.q0)
  sq0 = quat(0.0, s0)
  if abs(abs(q0s) - one(q0s)) == 0
    th = dual(th0, 0.5 * abs(quat(0, t)))
    th, dualquat(sq0)
  else
    th = dual(th0, 0.5 * dot(t, s0))
    s0c1 = cross(s0, t)
    tanth = tan(th0)
    s0c2 = (s0c1 / tanth + t) * 0.5
    sqev = cross(s0c2, s0)
    th, dualquat(sq0, quat(0.0, sqev))
  end
end

function angle(dq::DualQuaternion)
  th, ax = angleaxis(dq)
  th
end

function axis(dq::DualQuaternion)
  th, ax = angleaxis(dq)
  ax
end

function exp(dq::DualQuaternion)
  se = dual(dq.q0.s, dq.qe.s)
  se = exp(se)
  dq = dualquat(quat(0.0, imag_part(dq.q0)...), quat(0.0, imag_part(dq.qe)...))
  dq, th = normalizea(dq)
  if dq.norm
    dualquat(se) * (dualquat(cos(th)) + dq * dualquat(sin(th)))
  else
    dualquat(se)
  end
end

function log(dq::DualQuaternion)
  dq, a = normalizea(dq)
  sl = log(a)
  th, s = angleaxis(dq)
  s * dualquat(th) + dualquat(sl)
end

(^)(dq::DualQuaternion, dw::DualQuaternion) = exp(dw * log(dq))

function sqrt(dq::DualQuaternion)
  exp(0.5 * log(dq))
end

dualquatrand() = dualquat(quatrand(), quatrand())
ndualquatrand() = normalize(dualquatrand())

function rand(rng::AbstractRNG, ::Random.SamplerType{DualQuaternion{T}}) where {T<:Real}
    return DualQuaternion{T}(rand(rng, Quaternion{T}), rand(rng, Quaternion{T}), false)
end


#special type which represents an LUMatrix with permutations
struct LUMatrix{T}
    m::Matrix{T}
    p::Vector{Int64}
end

#another testing function that makes life much easier:
luupper(lu::LUMatrix{T}) where {T} = lu.m |> UpperTriangular
lulower(lu::LUMatrix{T}) where {T} = (lu.m - luupper(lu) + eye(T, size(lu.m,1))) |> LowerTriangular

#just a testing function that tests the LU matrix.
function lureassemble(lu::LUMatrix{T}) where {T}
    lulower(lu) * luupper(lu)
end

# TODO: implement unit testing on this function.

@generated function solve_quire(A::LowerTriangular{T}, B::AbstractVector{T}, quire = Quire(T), ::Type{Val{onediag}} = Val{false}) where {T, onediag}
    first_factor, row_factor = if onediag
        :(B[1]), :(tmp)
    else
        :(B[1] / A[1,1]), :(tmp / A[row, row])
    end

    quote
        m,n = size(A); l = length(B)
        (l == m) || throw(DimensionMismatch("matrix is ($m x $n) and vector is $l"))
        res = Vector{T}(n)
        tmp = zero(T)

        #start with the first row, which does not require a quire.
        @inbounds res[1] = $first_factor
        #TODO: reimplement second row with a proper fused-multiply subtract instead of
        #quire.
        for row = 2:n
            set!(quire, B[row])
            for col = 1:row
                @inbounds tmp =fdp!(quire, res[col], -A[row, col])
            end
            @inbounds res[row] = $row_factor
        end
        res
    end
end

function solve_quire(A::UpperTriangular{T}, B::AbstractVector{T}, quire = Quire(T)) where {T}
    m,n = size(A); l = length(B)
    (l == m) || throw(DimensionMismatch("matrix is ($m x $n) and vector is $l"))
    res = Vector{T}(n)
    tmp = zero(T)

    #start with the last row, which does not require a quire.
    @inbounds res[m] = B[m] / A[m,m]
    #TODO: reimplement penultimate row with a proper fused-multiply subtract instead of
    #quire.
    for row = (m-1):-1:1
        @inbounds set!(quire, B[row])
        for col = (row+1):m
            @inbounds tmp =fdp!(quire, res[col], -A[row, col])
        end
        @inbounds res[row] = tmp / A[row, row]
    end
    res
end

function solve_quire(LM::LUMatrix{T}, v::AbstractVector{T}, quire = Quire(T)) where {T}
    M = LM.m
    p = LM.p
    #first, solve the lower matrix against the permuted values in the vector.

    m,n = size(M); l = length(v)
    (l == m) || throw(DimensionMismatch("matrix is ($m x $n) and vector is $l"))
    res = Vector{T}(m)
    tmp = zero(T)

    #next, permute the incoming vector
    for idx = 1:m
        res[idx] = v[p[idx]]
    end

    res = solve_quire(LowerTriangular(M), res, quire, Val{true})
    #CONTINUE WITH CODE FOR THE UPPER MATRIX.

    solve_quire(UpperTriangular(M), res, quire)
end

function lu_factors!(M::AbstractMatrix{T}, quire = Quire(T))::LUMatrix{T} where {T}
    # lufact_quire is only going to work on square matrices at this time.
    # this returns an LU factorized A along with a pivot permutation.
    m, n = size(M)
    info = 0
    #set up the pivot list.
    pivot_list = collect(1:m)
    @inbounds begin
        #here we use the doolittle algorithm for LU factorization because it
        #permits the use of the quire for numerical accuracy.  The doolittle
        #algorithm operates on one row of the upper matrix, followed by the
        #corresponding column of the lower matrix, so activity should be index
        #ed by diagonal.
        for diag = 1:m
            # here we're going to do pivots to optimize the process.
            # search for the maximum absolute value in this column.
            pidx = indmax(abs.(M[diag:m, diag])) + diag - 1
            if pidx != diag
                #swap these rows
                (pivot_list[diag], pivot_list[pidx]) = (pivot_list[pidx], pivot_list[diag])

                #do the row interchange.
                if M[pidx, diag] != zero(T)
                    for col = 1:m
                        (M[diag, col], M[pidx, col]) = (M[pidx, col], M[diag, col])
                    end
                else
                    pivot_list[diag] = 0
                end
            end

            # solve columns in the upper matrix
            for upper_col = diag:m
                #signal which dot product we're gonna do.
                #then set up a
                accum = M[diag, upper_col]
                set!(quire, accum)
                for lower_col = 1:(diag-1)
                    accum = fdp!(
                      quire,
                      M[diag, lower_col],
                      -M[lower_col, upper_col])
                end
                M[diag, upper_col] = accum
            end

            #solve rows in the lower matrix
            factor = inv(M[diag, diag])
            for lower_row = (diag + 1):n
                accum = M[lower_row, diag]
                set!(quire, accum)
                for upper_row = 1:(diag-1)
                    accum = fdp!(
                      quire,
                      M[lower_row, upper_row],
                      -M[upper_row, diag])
                end
                M[lower_row, diag] = accum * factor
            end
        end
    end
    LUMatrix(M, pivot_list)
end

function solve_quire(A::AbstractMatrix{T}, B::AbstractVector{T}, quire = Quire(T)) where {T}
    m, n = size(A)
    if m == n
        if istril(A)
            if istriu(A)
                return Diagonal(A) \ B
            else
                return solve_quire(LowerTriangular(A), B, quire)
            end
        end
        if istriu(A)
            return solve_quire(UpperTriangular(A), B, quire)
        end
        return solve_quire(lu_factors!(copy(A)), B, quire)
    end
    throw(ArgumentError("currently, nonsquare matrices are not supported"))
end

function solve_quire_refine(A::AbstractMatrix{T}, B::AbstractVector{T}) where {T}
    res = A \ B
    res = refine(res, A, B)
    res = refine(res, A, B)
end


@testset "Pressure camera" begin

@testset "Parallel ray generator" begin
    ray = GRAYTR.ray_parallel(2.0, 0.0, 0.0)
    @test ray.origin ≈ Point3(0.0, 0.0, 1.0)
    @test ray.direction ≈ Vector3(0.0, 0.0, -1.0)
    @test norm(ray.direction) ≈ 1.0

    # Need to explicitly loosen test tolerance here because of float instability
    # with trinometry. There are some vector components that should be zero but get
    # values on the order of 1e-16.

    ray = GRAYTR.ray_parallel(2.0, 1.0, 0.0)
    @test ray.direction ≈ Vector3(0.0, 0.0, -1.0)
    @test isapprox(ray.origin, Point3(2.0, 0.0, 1.0) ; atol=1e-15)

    ray = GRAYTR.ray_parallel(2.0, 0.0, 1.0)
    @test ray.direction ≈ Vector3(0.0, 0.0, -1.0)
    @test ray.origin ≈ Point3(0.0, 0.0, 1.0)

    ray = GRAYTR.ray_parallel(2.0, 1.0, 1.0)
    @test ray.direction ≈ Vector3(0.0, 0.0, -1.0)
    @test isapprox(ray.origin, Point3(2.0, 0.0, 1.0) ; atol=1e-15)

    ray = GRAYTR.ray_parallel(2.0, 1.0, 0.5)
    @test isapprox(ray.origin, Point3(-2.0, 0.0, 1.0) ; atol=1e-15)

end

@testset "Point-source ray generator" begin
    ray = GRAYTR.ray_from_point(2.0, 4.0, 1.0, 0.0)
    @test ray.origin ≈ Point3(0.0, 0.0, 4.0)
    @test ray.direction ≈ Vector3(0.0, 0.0, -1.0)
    @test norm(ray.direction) ≈ 1.0

    # Need to explicitly loosen test tolerance here because of float instability
    # with trinometry. There are some vector components that should be zero but get
    # values on the order of 1e-16.

    ray = GRAYTR.ray_from_point(2.0, 4.0, 1.0, 1.0)
    @test ray.origin ≈ Point3(0.0, 0.0, 4.0)
    @test ray.direction ≈ Vector3(0.0, 0.0, -1.0)
    @test norm(ray.direction) ≈ 1.0

    ct = sqrt(1.0 - (2.0 / 4.0)^2)

    ray = GRAYTR.ray_from_point(2.0, 4.0, 0.0, 0.0)
    @test ray.origin ≈ Point3(0.0, 0.0, 4.0)
    @test isapprox(ray.direction, Vector3(sqrt(1-ct^2), 0.0, -ct), atol=1e-15)
    @test norm(ray.direction) ≈ 1.0

    ray = GRAYTR.ray_from_point(2.0, 4.0, 0.0, 0.5)
    @test ray.origin ≈ Point3(0.0, 0.0, 4.0)
    @test isapprox(ray.direction, Vector3(-sqrt(1-ct^2), 0.0, -ct), atol=1e-15)
    @test norm(ray.direction) ≈ 1.0

    ray = GRAYTR.ray_from_point(2.0, 4.0, 0.0, 0.25)
    @test ray.origin ≈ Point3(0.0, 0.0, 4.0)
    @test isapprox(ray.direction, Vector3(0.0, sqrt(1-ct^2), -ct), atol=1e-15)
    @test norm(ray.direction) ≈ 1.0

    ray = GRAYTR.ray_from_point(2.0, 4.0, 0.0, 0.75)
    @test ray.origin ≈ Point3(0.0, 0.0, 4.0)
    @test isapprox(ray.direction, Vector3(0.0, -sqrt(1-ct^2), -ct), atol=1e-15)
    @test norm(ray.direction) ≈ 1.0

end

end


@testset "asymptotic_block.jl" begin
    @testset "estimate, stderr, and z" begin
        kernel = (ExponentialKernel() ∘ ScaleTransform(0.1)) ⊗ WhiteKernel()
        for nclasses in (2, 10, 100), nsamples in (10, 50, 100)
            for blocksize in (2, 5, 10, 50)
                # blocksize may no be greater than number of samples
                blocksize < nsamples || continue

                # sample predictions and targets
                dist = Dirichlet(nclasses, 1)
                predictions = [rand(dist) for _ in 1:nsamples]
                targets_consistent = [
                    rand(Categorical(prediction)) for prediction in predictions
                ]
                targets_onlyone = ones(Int, length(predictions))

                skce = BlockUnbiasedSKCE(kernel, blocksize)

                # for both sets of targets
                for targets in (targets_consistent, targets_onlyone)
                    test = AsymptoticBlockSKCETest(kernel, blocksize, predictions, targets)

                    @test test.blocksize == blocksize
                    @test test.nblocks == nsamples ÷ blocksize
                    @test test.estimate ≈ skce(predictions, targets)
                    @test test.z == test.estimate / test.stderr

                    @test pvalue(test) ==
                          pvalue(test; tail=:right) ==
                          ccdf(Normal(), test.z)
                    @test_throws ArgumentError pvalue(test; tail=:left)
                    @test_throws ArgumentError pvalue(test; tail=:both)

                    for α in 0.55:0.05:0.95
                        q = quantile(Normal(), α)
                        @test confint(test; level=α) ==
                              confint(test; level=α, tail=:right) ==
                              (max(0, test.estimate - q * test.stderr), Inf)
                        @test_throws ArgumentError confint(test; level=α, tail=:left)
                        @test_throws ArgumentError confint(test; level=α, tail=:both)
                    end
                end
            end
        end
    end

    @testset "consistency" begin
        kernel1 = ExponentialKernel() ∘ ScaleTransform(0.1)
        kernel2 = WhiteKernel()
        kernel = kernel1 ⊗ kernel2
        αs = 0.05:0.1:0.95
        nsamples = 100

        pvalues_consistent = Vector{Float64}(undef, 100)
        pvalues_onlyone = similar(pvalues_consistent)

        for blocksize in (2, 5, 10)
            for nclasses in (2, 10)
                dist = Dirichlet(nclasses, 1)
                predictions = [Vector{Float64}(undef, nclasses) for _ in 1:nsamples]
                targets_consistent = Vector{Int}(undef, nsamples)
                targets_onlyone = ones(Int, nsamples)

                for i in eachindex(pvalues_consistent)
                    # sample predictions and targets
                    for j in 1:nsamples
                        rand!(dist, predictions[j])
                        targets_consistent[j] = rand(Categorical(predictions[j]))
                    end

                    # define test
                    test_consistent = AsymptoticBlockSKCETest(
                        kernel, blocksize, predictions, targets_consistent
                    )
                    test_onlyone = AsymptoticBlockSKCETest(
                        kernel, blocksize, predictions, targets_onlyone
                    )

                    # estimate pvalues
                    pvalues_consistent[i] = pvalue(test_consistent)
                    pvalues_onlyone[i] = pvalue(test_onlyone)
                end

                # compute empirical test errors
                ecdf_consistent = ecdf(pvalues_consistent)
                @test all(ecdf_consistent(α) ≤ α + 0.1 for α in αs)
                @test all(p < 0.05 for p in pvalues_onlyone)
            end
        end
    end
end


@testset "dotproduct" begin
    A = rand(10, 5)
    B = rand(20, 5)
    d = KernelFunctions.DotProduct()
    @test diag(pairwise(d, A; dims=2)) == [dot(A[:, i], A[:, i]) for i in 1:size(A, 2)]
    @test_throws DimensionMismatch d(rand(3), rand(4))
    @test d(3.0, 2.0) == 6.0
end


module VectorAutoregressions

using DataFrames, TimeSeries
using Distributions: ccdf, TDist, Chisq, FDist
using Dates, LinearAlgebra
using StatsBase: CoefTable, autocov, autocor, crosscov, crosscor

import Base: show
import StatsBase: loglikelihood, aic, sample, coef, stderror, residuals, coeftable, fitted

export VAREstimate,
    VAR,
    LagSelection,
    StabilityCheck,
    PortmanteauTest,
    LMCorrTest,
    loglikelihood,
    aic,
    sic,
    hqc,
    coef,
    stderror,
    residuals,
    coeftable,
    fitted

include("matrix_utilities.jl")
include("VAR.jl")
include("statsbase.jl")
include("diagnostics.jl")
include("lag_selection.jl")
include("IRF.jl")

end # module


module MagnaGA

using SparseArrays
using Random
using ..GAFramework
import ..GAFramework: fitness, crossover!, mutation!, selection, randcreature, genauxga, printfitness
using ..CayleyCrossover
using ..Spmap

export MagnaCreature, MagnaModel

# Represents a permutation, i.e. 1-1 mapping from one set to another
struct MagnaCreature
    f :: Vector{Int}
    edge_score :: Float64
    node_score :: Float64
    score :: Float64 # alignment score
end

fitness(x::MagnaCreature) where {T} = x.score

"""
Model to find network alignment by optimizing S3 (see MAGNA paper)
A network alignment is represented using a permutation
Faster implementation, slightly slower than the C++ code from the MAGNA++ paper
    using GAFramework, GAFramework.MagnaGA
    using LightGraphs
    using Random
    mv = 100
    me = 400
    G1 = adjacency_matrix(LightGraphs.erdos_renyi(mv,me))
    perm = randperm(mv)
    G2 = G1[invperm(perm),invperm(perm)]
    # G1.|V| <= G2.|V|
    S = zeros(mv,mv)
    alpha = 1.0
    model = MagnaModel(G1,G2,S,alpha)
    st = GAState(model, ngen=100, npop=6_000, elite_fraction=0.1,
        params=Dict(:print_fitness_iter=>1))
    ga!(st)
    println("accuracy = ", sum(st.pop[1].f .== perm)/mv)
    st.pop[1]
"""
struct MagnaModel <: GAModel
    G1   :: SparseMatrixCSC{Int,Int}
    G2   :: SparseMatrixCSC{Int,Int}
    S :: Matrix{Float64}
    alpha :: Float64
    function MagnaModel(G1::SparseMatrixCSC,G2::SparseMatrixCSC,S::Matrix,alpha::Number)
        (size(G1,1)==size(G1,2) && size(G2,1)==size(G2,1) &&
            size(S,1)==size(G1,1) && size(S,2)==size(G2,1)) || error("input sizes")
        new(G1,G2,S,alpha)
    end
end

"""
    using GAFramework, GAFramework.MagnaGA, GAFramework.PermGA
    using LightGraphs, SparseArrays
    using Random
    GAFramework.setthreads(false)
    nv = 100
    ne = 400
    G2 = adjacency_matrix(LightGraphs.erdos_renyi(nv,ne))
    perm = randperm(nv)
    mv = 75
    G1 = G2[perm[1:mv], perm[1:mv]]
    S = rand(Float64, mv, nv)/10
    for i in 1:mv
        S[i,i] = 0.95
    end
    alpha = 0.5
    model1 = MagnaModel(G1,G2,S,alpha)
    model2 = NetalignModel(G1,G2)
    aux1 = genauxga(model1)
    f = randperm(size(G2,1));
    c1 = MagnaCreature(f, model1, aux1)
    c2 = PermCreature(f, model2)
    c1.score, c2.score

    st1 = GAState(model1, ngen=10, npop=60_000, elite_fraction=0.1, rng=MersenneTwister(10));
    st2 = GAState(model2, ngen=10, npop=60_000, elite_fraction=0.1, rng=MersenneTwister(10));

    @time ga!(st1)
    @time ga!(st2)

"""

function genauxga(m::MagnaModel)
    crossover_aux = randperm(size(m.G2,1)), fill(false, size(m.G2,1))
    Spmap.map_index_aux(m.G1, m.G2), crossover_aux
end

# alpha * S_E + (1-alpha) * S_N
function MagnaCreature(h::AbstractVector, m::MagnaModel, aux)
    length(h)==size(m.G2,1) || error("length of h")
    # aux = genauxga(..) can be used to reduce allocations here
    # by using auxiliary space instead of reallocating
    ((K, hinv, hinvrowvals), _) = aux
    # Assumes that all edge weights are 1
    Spmap.plus_getindex!(K, m.G1, m.G2, h, hinv, hinvrowvals) # K = G1 + G2[h,h][1:m,1:m]

    w = nonzeros(K)
    Nc = count(x->x==2, w)
    Nn = count(x->x==1, w)

    Nc,cr = divrem(Nc,2)
    Nn,nr = divrem(Nn,2) # need this & above so it works well with sim. ann. code
    if cr!=0 || nr!=0 error("G1 and G2 need to be symmetric"); end
    edge_score = Nc/(Nc + Nn)

    node_score = 0.0
    if m.alpha != 1.0
        S = m.S
        @inbounds for i in 1:size(m.G1,1)
            node_score += S[i, h[i]]
        end
        node_score /= size(m.G1,1)
    end

    score = m.alpha*edge_score + (1.0-m.alpha) * node_score

    MagnaCreature(h, edge_score, node_score, score)
end

function randcreature(m::MagnaModel, aux, rng::AbstractRNG) where {T}
    f = randperm(rng, size(m.G2,1))
    MagnaCreature(f, m, aux)
end

function crossover!(z, x, y, st::GAState{MagnaModel}, aux, rng::AbstractRNG) where {T}
    (_, (r, visited)) = aux
    visited .= false
    cayley_crossover!(z.f, x.f, y.f, rng, r, visited)
    MagnaCreature(z.f, st.model, aux)
end

function printfitness(curgen::Int, x::MagnaCreature)
    println("curgen: $(curgen) edge_score: $(x.edge_score) node_score: $(x.node_score) score: $(x.score)")
end

end # MagnaGA


using Test
using KernelFunctions
using Random: seed!

seed!(42)

dims = (10,5)
X = rand(dims...)
x = rand(dims[1])
s = 3.0
v1 = vcat(3.0,4.0*ones(dims[2]-1))
v2 = vcat(3.0,4.0*ones(dims[1]-1))
P = rand(5,10)
sdims = [1,2,3]
f(x) = sin.(x)

@testset "Transform Test" begin
    ## Test Scale Transform
    @testset "IdentityTransform" begin
        @test KernelFunctions.apply(IdentityTransform(),X)==X
    end
    @testset "ScaleTransform" begin
        t = ScaleTransform(s)
        @test all(KernelFunctions.apply(t,X).==s*X)
        s2 = 2.0
        KernelFunctions.set!(t,s2)
        @test all(t.s.==[s2])
        @test isequal(ScaleTransform(s),ScaleTransform(s))
    end
    ## Test ARD Transform
    @testset "ARDTransform" begin
        vt1 = ARDTransform(v1)
        vt2 = ARDTransform(v2)
        @test all(KernelFunctions.apply(vt1,X,obsdim=1).==v1'.*X)
        @test all(KernelFunctions.apply(vt2,X,obsdim=2).==v2.*X)
        newv1 = rand(5)
        KernelFunctions.set!(vt1,newv1)
        @test all(vt1.v .== newv1)
        @test ARDTransform(s,dims[2]).v == ARDTransform(s*ones(dims[2])).v
        @test_throws DimensionMismatch KernelFunctions.apply(vt1,rand(3,4))
    end
    ## Test LowRankTransform
    @testset "LowRankTransform" begin
        tp = LowRankTransform(P)
        @test all(KernelFunctions.apply(tp,X,obsdim=2).==P*X)
        @test all(KernelFunctions.apply(tp,x).==P*x)
        @test tp.proj == P
        P2 = rand(5,10)
        KernelFunctions.set!(tp,P2)
        @test all(tp.proj.==P2)
        @test_throws AssertionError KernelFunctions.set!(tp,rand(6,10))
        @test_throws DimensionMismatch KernelFunctions.apply(tp,rand(11,3))
    end
    ## Test FunctionTransform
    @testset "FunctionTransform" begin
        tf = FunctionTransform(f)
        KernelFunctions.apply(tf,X,obsdim=1)
        @test all(KernelFunctions.apply(tf,X,obsdim=1).==f(X))
    end
    ## Test SelectTransform
    @testset "SelectTransform" begin
        ts = SelectTransform(sdims)
        @test all(KernelFunctions.apply(ts,X,obsdim=2).==X[sdims,:])
        @test all(KernelFunctions.apply(ts,x).==x[sdims])
        sdims2 = [2,3,5]
        KernelFunctions.set!(ts,sdims2)
        @test all(ts.select.==sdims2)
    end
    ## Test ChainTransform
    @testset "ChainTransform" begin
        t = ScaleTransform(s)
        tp = LowRankTransform(P)
        tf = FunctionTransform(f)
        tchain = ChainTransform([t,tp,tf])
        @test all(KernelFunctions.apply(tchain,X,obsdim=2).==f(P*(s*X)))
        @test all(KernelFunctions.apply(tchain,X,obsdim=2).==
                    KernelFunctions.apply(tf∘tp∘t,X,obsdim=2))
    end
end


module RulesBetaOutTest

using Test
using ReactiveMP
using Random
import ReactiveMP: @test_rules

@testset "rules:Beta:out" begin

     @testset "Belief Propagation: (m_a::PointMass, m_b::PointMass)" begin

        @test_rules [ with_float_conversions = true ] Beta(:out, Marginalisation) [
            (input = (m_a = PointMass(1.0), m_b = PointMass(2.0)), output = Beta(1.0, 2.0)),
            (input = (m_a = PointMass(2.0), m_b = PointMass(2.0)), output = Beta(2.0, 2.0)),
            (input = (m_a = PointMass(3.0), m_b = PointMass(3.0)), output = Beta(3.0, 3.0)),
        ]

    end

    @testset "Variational Message Passing: (q_a::PointMass, q_b::PointMass)" begin

        @test_rules [ with_float_conversions = true ] Beta(:out, Marginalisation) [
            (input = (q_a = PointMass(1.0), q_b = PointMass(2.0)), output = Beta(1.0, 2.0)),
            (input = (q_a = PointMass(2.0), q_b = PointMass(2.0)), output = Beta(2.0, 2.0)),
            (input = (q_a = PointMass(3.0), q_b = PointMass(3.0)), output = Beta(3.0, 3.0)),
        ]

    end
end
end


n = 3
m = 2

for elty in (Float32, Float64, ComplexF32, ComplexF64)

    A0 = randn(elty, n, n)
    B0 = randn(elty, n, m)
    B1 = randn(elty, n, m)

    A = copy(A0)

    # Full matrix
    linws = FastLapackInterface.LinSolveWs{elty,Int64}(n)

    FastLapackInterface.lu!(A, linws)
    @test A == A0
    F = lu(A)
    @test UpperTriangular(reshape(linws.lu, n, n)) ≈ F.U

    FastLapackInterface.lu!(A', linws)
    @test UpperTriangular(reshape(linws.lu, n, n)) ≈ F.U

    B = copy(B0)
    FastLapackInterface.linsolve_core!(A, B, linws)
    @test A == A0
    @test B ≈ A \ B0

    copy!(B, B1)
    FastLapackInterface.linsolve_core_no_lu!(A, B, linws)
    @test A == A0
    @test B ≈ A \ B1

    copy!(B, B0)
    FastLapackInterface.linsolve_core!(A', B, linws)
    @test B ≈ A' \ B0

    copy!(B, B1)
    FastLapackInterface.linsolve_core_no_lu!(A', B, linws)
    @test A == A0
    @test B ≈ A' \ B1

    # view of a matrix in upper left corner
    linws1 = FastLapackInterface.LinSolveWs{elty,Int64}(n - 1)
    C = view(A, 1:n-1, 1:n-1)
    D = view(B, 1:n-1, 1:m-1)
    D0 = copy(D)
    D1 = copy(D)
    FastLapackInterface.linsolve_core!(C, D, linws1)
    @test C == view(A, 1:n-1, 1:n-1)
    @test D ≈ C \ D0

    D = copy(D1)
    FastLapackInterface.linsolve_core_no_lu!(C, D, linws1)
    @test C == view(A, 1:n-1, 1:n-1)
    @test D ≈ C \ D1

    D = copy(D0)
    FastLapackInterface.linsolve_core!(C', D, linws1)
    @test C == view(A, 1:n-1, 1:n-1)
    @test D ≈ C' \ D0

    D = copy(D1)
    FastLapackInterface.linsolve_core_no_lu!(C', D, linws1)
    @test C == view(A, 1:n-1, 1:n-1)
    @test D ≈ C' \ D1

    # view of a matrix in lower left corner
    linws1 = FastLapackInterface.LinSolveWs{elty,Int64}(n - 1)
    C = view(A, 2:n, 1:n-1)
    C1 = copy(C)
    FastLapackInterface.lu!(C1, linws1)
    F = LinearAlgebra.lu!(C)
    @test triu(reshape(linws1.lu, n - 1, n - 1)) ≈ F.U
    @test tril(reshape(linws1.lu, n - 1, n - 1), -1) ≈ tril(F.L, -1)

    D0 = view(B0, 2:n, 1:m-1)
    D1 = view(B1, 2:n, 1:m-1)
    D = copy(D0)
    FastLapackInterface.linsolve_core!(C, D, linws1)
    @test C == view(A, 2:n, 1:n-1)
    @test D ≈ C \ D0

    D = copy(D1)
    FastLapackInterface.linsolve_core_no_lu!(C, D, linws1)
    @test C == view(A, 2:n, 1:n-1)
    @test D ≈ C \ D1

    D = copy(D0)
    FastLapackInterface.linsolve_core!(C', D, linws1)
    @test C == view(A, 2:n, 1:n-1)
    @test D ≈ C' \ D0

    D = copy(D1)
    FastLapackInterface.linsolve_core_no_lu!(C', D, linws1)
    @test C == view(A, 2:n, 1:n-1)
    @test D ≈ C' \ D1

    # using too big a work space
    linws1 = FastLapackInterface.LinSolveWs{elty,Int64}(n)
    C = view(A, 2:n, 1:n-1)
    C1 = copy(C)
    FastLapackInterface.lu!(C1, linws1)
    F = LinearAlgebra.lu!(C)
    @test triu(reshape(linws1.lu[1:(n-1)^2], n - 1, n - 1)) ≈ F.U
    @test tril(reshape(linws1.lu[1:(n-1)^2], n - 1, n - 1), -1) ≈ tril(F.L, -1)

    D0 = view(B0, 2:n, 1:m-1)
    D1 = view(B1, 2:n, 1:m-1)
    D = copy(D0)
    FastLapackInterface.linsolve_core!(C, D, linws1)
    @test C == view(A, 2:n, 1:n-1)
    @test D ≈ C \ D0

    D = copy(D1)
    FastLapackInterface.linsolve_core_no_lu!(C, D, linws1)
    @test C == view(A, 2:n, 1:n-1)
    @test D ≈ C \ D1

    D = copy(D0)
    FastLapackInterface.linsolve_core!(C', D, linws1)
    @test C == view(A, 2:n, 1:n-1)
    @test D ≈ C' \ D0

    D = copy(D1)
    FastLapackInterface.linsolve_core_no_lu!(C', D, linws1)
    @test C == view(A, 2:n, 1:n-1)
    @test D ≈ C' \ D1

end


using PyPlot, transD_GP, Random, LinearAlgebra, Statistics

## model
Random.seed!(3)
nlayers, dz = 50, 5.
ρbg        = 1000
zfixed     = [-1e5,   (0:dz:(nlayers-1)*dz )...]
rho        = [1e12,   10 .^(log10(ρbg) .+ randn(nlayers))...]
ianom      = 5
rho[ianom] = 10.
δ          = 1e-4  # fd step size in log10

##  geometry and frequencies for VMD in frequency domain
nFreqsPerDecade     = 5
freqLowLimit        = 1e-3
freqHighLimit       = 1e5
freqs = 10 .^(log10(freqLowLimit):1/nFreqsPerDecade:log10(freqHighLimit))
rRx = 100.
zRx = -37.
zTx = -35.
nkᵣeval = 50
modelprimary = true

Fvmd = transD_GP.AEM_VMD_HMD.HFieldDHT(
                      calcjacobian = true,
                      nmax = length(zfixed),
                      zTx    = zTx,
                      rRx    = rRx,
                      freqs  = freqs,
                      zRx    = zRx,
                      nkᵣeval = nkᵣeval,
                      modelprimary = modelprimary)
transD_GP.AEM_VMD_HMD.getfieldFD!(Fvmd, zfixed, rho)

## plot
f, ax = plt.subplots(1,3, sharey=true, figsize=(8,6))
delz, delf = diff(zfixed[2:end]), diff(log10.(freqs))
flabel = [(log10.(freqs) .- delf[1]/2)..., log10.(freqs)[end]+delf[end]/2]
zlab = [zfixed[2:end]..., zfixed[end]+delz[end]]
ax[1].step(log10.(rho[2:end]), zfixed[2:end])
im2 = ax[2].pcolormesh(flabel, zlab, abs.(Fvmd.HFD_z_J[2:end,:]), shading="auto")
cb2 = colorbar(im2, ax=ax[2])
im3 = ax[3].pcolormesh(flabel, zlab, log10.(abs.(Fvmd.HFD_z_J[2:end,:])))
cb2 = colorbar(im3, ax=ax[3])
ax[1].set_xlabel("log10 Ω-m")
ax[1].invert_xaxis()
ax[2].invert_yaxis()
ax[2].set_xlabel("log10 Hz")
ax[3].set_xlabel("log10 Hz")
ax[1].set_ylabel("Depth m")
ax[2].set_title("∂f/∂(log10σ)")
ax[3].set_title("∂f/∂(log10σ)")

## approximate partials with central differences and compare
J = copy(Fvmd.HFD_z_J)

function getJ!(Fin, zin, ρ, Δ)
    # sensitivities computed in log10σ
    # same as in AEM_VMD_HMD code
    Jout = similar(Fin.HFD_z_J)
    modelfd = copy(ρ)
    jflag = Fin.calcjacobian
    Fin.calcjacobian = false
    for i = 2:length(zin)
        modelfd[:] = ρ
        modelfd[i] = 10 .^-(-log10(ρ[i]) .+ Δ/2)
        transD_GP.AEM_VMD_HMD.getfieldFD!(Fin, zin, modelfd)
        HFD_z1 = copy(Fin.HFD_z)
        modelfd[i] = 10 .^-(-log10(ρ[i]) .- Δ/2)
        transD_GP.AEM_VMD_HMD.getfieldFD!(Fin, zin, modelfd)
        Jout[i,:] = (HFD_z1 - Fin.HFD_z) / Δ
    end
    Fin.calcjacobian = jflag
    Jout
end
J_fd = getJ!(Fvmd, zfixed, rho, δ)

## comparison plots
f, ax = plt.subplots(1,3, sharey=true, figsize=(10,6))
ax[1].step(log10.(rho[2:end]), zfixed[2:end])
im2 = ax[2].pcolormesh(flabel, zlab, abs.(J[2:end,:]))
cb2 = colorbar(im2, ax=ax[2])
im3 = ax[3].pcolormesh(flabel, zlab, abs.(J_fd[2:end,:]))
cb2 = colorbar(im3, ax=ax[3])
ax[1].set_xlabel("log10 Ω-m")
ax[1].invert_xaxis()
ax[2].invert_yaxis()
ax[2].set_xlabel("log10 Hz")
ax[3].set_xlabel("log10 Hz")
ax[1].set_ylabel("Depth m")
ax[2].set_title("∂f/∂(log10σ) analytic")
ax[3].set_title("∂f/∂(log10σ) FD")

## more detailed
figure()
nlayer = 3
subplot(121)
semilogx(freqs, real(J_fd[nlayer,:]), "--k")
semilogx(freqs, real(J[nlayer,:]), "-")
subplot(122)
semilogx(freqs, imag(J_fd[nlayer,:]), "--k")
semilogx(freqs, imag(J[nlayer,:]), "-")


################################################################################
#
#   fasta : Accelerated Forward Backward splitting with
#            adaptative step size and backtracking
#
################################################################################

const NLookback=10 # previous steps lookback to decide backtracking

function fasta!{T<:FloatingPoint}(
    rmodel::SupervisedRiskModel,     # model definition
    θ::StridedArray{T},              # initial parameter guess
    x::StridedArray{T},              # inputs
    y::StridedArray{T};              # responses
    reg::Regularizer = ZeroReg(),    # regularizer
    λ::T             = one(T),       # initial step size
    cbinterval::Int  = 0,            # callback interval
    callback         = fasta_trace,  # callback function
    maxsteps::Int    = 100,          # max number of steps
    maxtime::Int     = 0         )   # max calculation time (0 = no limit)

    g              = similar(θ) # holds gradient
    prevg          = similar(θ) # gradient of previous step
    deltag         = similar(θ) # g - predg

    prevθ          = similar(θ) # θ of previous step
    deltaθ         = similar(θ) # θ - prevθ

    dualθ          = copy(θ)    # dual θ
    candidatedualθ = similar(θ) # candidate dual θ
    deltadualθ     = similar(θ)

    vhist     = fill(-Inf, NLookback) # past values

    function fbstep!(θ1, g, θ0, λ)
        copy!(θ1, θ0)
        axpy!(-λ, g, θ1)
        prox!(reg, θ1, θ1, λ)
    end

    ### Loop initialization
    starttime = time()
    α, μ      = 1., 0.      # acceleration parameters
    v, _      = value_and_addgrad!(rmodel, 0., g, 1., dualθ, x, y)
    vhist[1]  = v

    cbinterval > 0 && callback(0, θ, v, value(reg, θ), λ, α, 0)
    for t in 1:maxsteps
        (maxtime != 0) && (time() - starttime > maxtime) && break

        # acceleration
        if t >= 2
            if dot(g, deltadualθ) >= 0.
                α = 1.
                μ = 0.
            else
                candα = (1. + sqrt(1+4α*α))/2.
                μ = (α - 1.) / candα
                α = candα
            end
        end

        # adaptative step size
        if t >= 2
            copy!(deltag, g) ; axpy!(-1., prevg, deltag)
            tmp = dot(deltadualθ, deltag)
            λs  = dot(deltadualθ, deltadualθ) / tmp
            λm  = tmp / dot(deltag, deltag)
            if λs > 0. && λm > 0.
                λ = (λm / λs > 0.5) ? λm : λs - 0.5λm
            end
        end

        # update, with backtracking
        btsteps = 0
        maxv = maximum(vhist)
        copy!(prevθ, θ)
        while btsteps < 10
            fbstep!(θ, g, dualθ, λ)
            copy!(deltaθ, θ) ; axpy!(-1., prevθ, deltaθ)
            copy!(candidatedualθ, θ)
            axpy!(μ, deltaθ, candidatedualθ)

            candidatev = value(rmodel, candidatedualθ, x, y)
            deltadualθ = candidatedualθ - dualθ # TODO : preallocate
            if candidatev < maxv + dot(deltadualθ, g) + sumabs2(deltadualθ)/2λ
                break
            end
            λ /= 2.
            btsteps += 1
        end
        copy!(dualθ     , candidatedualθ)

        # update variables for next step
        copy!(prevg     , g)
        v, _ = value_and_addgrad!(rmodel, 0., g, 1., dualθ, x, y)
        unshift!(vhist, v)
        pop!(vhist)

        # callback
        if cbinterval > 0 && t % cbinterval == 0
            callback(t, θ, v, value(reg, θ), λ, α, btsteps)
        end
    end

    θ
end

# Non mutating version
function fasta{T<:FloatingPoint}(
    rmodel::SupervisedRiskModel,     # model definition
    θ::StridedArray{T},              # initial parameter guess
    x::StridedArray{T},              # inputs
    y::StridedArray{T};              # responses
    vargs... )

    fasta!(rmodel, copy(θ), x, y; vargs...)
end


# import libraries
using GeometricIntegrators
using ProgressMeter

# import Particles package
using Particles

# import HDF5 package for storing solution
using HDF5

# parameters
Δt = 0.1    # time step size
nt = 100    # number of time steps
ns = 5      # number of samples

# output file
h5file = "charged_particles.hdf5"

# random initial conditions
z₀ = collect(hcat(rand(ns) .* 2 .- 1, rand(ns) .* 2 .- 1, zeros(ns), rand(ns) .- 0.5, rand(ns) .- 0.5, ones(ns))')

# electric field
E(x::Vector{DT}) where {DT} = DT[0, 0, cos(2π*x[3])]

# magnetic field
B(x::Vector{DT}) where {DT} = DT[0, 0, 1]

# vector field
function lorentz_force!(t, z, ż)
    for i in axes(z,2)
        x = z[1:3,i]
        v = z[4:6,i]

        e = E(x)
        b = B(x)

        ż[1,i] = v[1]
        ż[2,i] = v[2]
        ż[3,i] = v[3]

        ż[4,i] = e[1] + v[2] * b[3] - v[3] * b[2]
        ż[5,i] = e[2] + v[3] * b[1] - v[1] * b[3]
        ż[6,i] = e[3] + v[1] * b[2] - v[2] * b[1]
    end
end

# solution storage
function copy_to_hdf5(h5z, z, n)
    h5z[:,:,n+1] = z
end

# create an Equation instance
equ = ODE(lorentz_force!, z₀)

# create HDF5 file and copy initial conditions
h5  = h5open(h5file, "w")
h5z = create_dataset(h5, "z", eltype(z₀), ((6, ns, nt+1), (6, ns, -1)), chunk=(6,ns,1))
copy_to_hdf5(h5z, z₀, 0)

# create integrator
int = IntegratorExplicitEuler(equ, Δt)

# create atomic solution
asol = AtomicSolution(equ)

# copy initial conditons to atomic solution
set_initial_conditions!(asol, equ)

# initilize integrator
initialize!(int, asol)

# loop over time steps showing progress bar
try
    @showprogress 5 for n in 1:nt
        integrate_step!(int, asol)
        copy_to_hdf5(h5z, asol.q, n)
    end
finally
    # close HDF5 file
    close(h5)
end

# load Plots package
using Plots

# read array from HDF5 file
z = h5read(h5file, "z")

# set plot ranges
xlim = (-2, +2)
ylim = (-2, +2)
zlim = ( 0, 20)

# create empty 3d plot with ns empty series
plt = plot3d(
    ns,
    xlim = xlim,
    ylim = ylim,
    zlim = zlim,
    title = "Charged Particles",
    legend = false,
    marker = 5,
    size = (800, 600)
)

# plot solutions
for j in axes(z,3)
    for i in axes(z,2)
        push!(plt, i, z[1,i,j], z[2,i,j], z[3,i,j])
    end
end

# save figure to file
savefig("charged_particles.png")


# FIELDS:
#==============================================================================#
"""
    abstract type Field end

abstract type at the top of the type hierarchy for denoting fields over which
vector spaces can be defined. Two common fields are `ℝ` and `ℂ`, representing the
field of real or complex numbers respectively.
"""
abstract type Field end

struct RealNumbers <: Field end
struct ComplexNumbers <: Field end

const ℝ = RealNumbers()
const ℂ = ComplexNumbers()

Base.show(io::IO, ::RealNumbers) = print(io,"ℝ")
Base.show(io::IO, ::ComplexNumbers) = print(io,"ℂ")

Base.in(::Any, ::Field) = false
Base.in(::Real, ::RealNumbers) = true
Base.in(::Number, ::ComplexNumbers) = true

Base.@pure Base.issubset(::Type, ::Field) = false
Base.@pure Base.issubset(::Type{<:Real}, ::RealNumbers) = true
Base.@pure Base.issubset(::Type{<:Number}, ::ComplexNumbers) = true
Base.@pure Base.issubset(::RealNumbers, ::ComplexNumbers) = true

# VECTOR SPACES:
#==============================================================================#
"""
    abstract type VectorSpace end

abstract type at the top of the type hierarchy for denoting vector spaces
"""
abstract type VectorSpace end

"""
    function fieldtype(V::VectorSpace) -> Field

Returns the field type over which a vector space is defined.
"""
function fieldtype end
fieldtype(V::VectorSpace) = fieldtype(typeof(V))

# Basic vector space methods
#----------------------------
"""
    space(a) -> VectorSpace

Returns the vector space associated to object `a`.
"""
function space end

"""
    dim(V::VectorSpace) -> Int
Returns the total dimension of the vector space `V` as an Int.
"""
function dim end

"""
    dual(V::VectorSpace) -> VectorSpace
Returns the dual space of `V`; also obtained via `V'`. It is assumed that
`typeof(V) == typeof(V')`.
"""
function dual end

# convenience definitions:
Base.adjoint(V::VectorSpace) = dual(V)
Base.:*(V1::VectorSpace, V2::VectorSpace) = ⊗(V1, V2)

# Hierarchy of elementary vector spaces
#---------------------------------------
"""
    abstract type ElementarySpace{k} <: VectorSpace end

Elementary finite-dimensional vector space over a field `k` that can be used as the
index space corresponding to the indices of a tensor. Every elementary vector space
should respond to the methods `conj` and `dual`, returning the complex conjugate space
and the dual space respectively. The complex conjugate of the dual space is obtained
as `dual(conj(V)) === conj(dual(V))`. These different spaces should be of the same type,
so that a tensor can be defined as an element of a homogeneous tensor product
of these spaces.
"""
abstract type ElementarySpace{k} <: VectorSpace end
const IndexSpace = ElementarySpace

fieldtype(::Type{<:ElementarySpace{k}}) where {k} = k

"""
    conj(V::ElementarySpace) -> ElementarySpace
Returns the conjugate space of `V`. For `fieldtype(V)==ℝ`, `conj(V) == V`
It is assumed that `typeof(V) == typeof(conj(V))`.
"""
Base.conj(V::ElementarySpace{ℝ}) = V

"""
    abstract type InnerProductSpace{k} <: ElementarySpace{k} end

abstract type for denoting vector with an inner product and a corresponding
metric, which can be used to raise or lower indices of tensors
"""
abstract type InnerProductSpace{k} <: ElementarySpace{k} end

"""
    abstract type EuclideanSpace{k<:Union{ℝ,ℂ}} <: InnerProductSpace{k} end

abstract type for denoting real or complex spaces with a standard (Euclidean)
inner product (i.e. orthonormal basis), such that the dual space is naturally
isomorphic to the conjugate space (in the complex case) or even to the
space itself (in the real case)
"""
abstract type EuclideanSpace{k} <: InnerProductSpace{k} end # k should be ℝ or ℂ

dual(V::EuclideanSpace) = conj(V)
# dual space is naturally isomorphic to conjugate space for inner product spaces

# representation spaces: we restrict to complex Euclidean space supporting unitary representations
"""
    abstract type AbstractRepresentationSpace{G<:Sector} <: EuclideanSpace{ℂ} end

Complex Euclidean space with a direct sum structure corresponding to different
different superselection sectors of type `G<:Sector`, e.g. the elements or
irreps of a compact or finite group, or the labels of a unitary fusion category.
"""
abstract type AbstractRepresentationSpace{G<:Sector} <: EuclideanSpace{ℂ} end

"""
    function sectortype(a) -> Sector

Returns the type of sector over which object `a` (e.g. a representation space or
an invariant tensor) is defined. Also works in type domain.
"""
sectortype(V::VectorSpace) = sectortype(typeof(V))
sectortype(::Type{<:ElementarySpace}) = Trivial
sectortype(::Type{<:AbstractRepresentationSpace{G}}) where {G} = G

"""
    function sectors(a)

    Returns the different sectors of object `a`( e.g. a representation space or
    an invariant tensor).
"""
sectors(::ElementarySpace) = (Trivial(),)
dim(V::ElementarySpace, ::Trivial) = sectortype(V) == Type{Trivial} ? dim(V) : throw(SectorError())

# Composite vector spaces
#-------------------------
"""
    abstract type CompositeSpace{S<:ElementarySpace} <: VectorSpace end

Abstract type for composite spaces that are defined in terms of a number of
elementary vector spaces of a homogeneous type `S<:ElementarySpace{k}`.
"""
abstract type CompositeSpace{S<:ElementarySpace} <: VectorSpace end

fieldtype(::Type{<:CompositeSpace{S}}) where {S<:ElementarySpace} = fieldtype(S)
sectortype(::Type{<:CompositeSpace{S}}) where {S<:ElementarySpace} = sectortype(S)

# Specific realizations of ElementarySpace types
#------------------------------------------------
# spaces without internal structure
include("cartesianspace.jl")
include("complexspace.jl")
include("generalspace.jl")
include("representationspace.jl")

# # Specific realizations of CompositeSpace types
# #-----------------------------------------------
include("productspace.jl")
# include("superspace.jl")
# include("invariantspace.jl")

# Other examples might include:
# braidedspace and fermionspace
# symmetric and antisymmetric subspace of a tensor product of identical vector spaces
# ...