using CUDA.CUSOLVER

using LinearAlgebra
using LinearAlgebra: BlasInt

m = 15
n = 10
l = 13
k = 1

@test_throws ArgumentError CUSOLVER.cusolverjob('M')

@testset "elty = $elty" for elty in [Float32, Float64, ComplexF32, ComplexF64]
    @testset "Cholesky (po)" begin
        A    = rand(elty,n,n)
        A    = A*A' #posdef
        B    = rand(elty,n,n)
        d_A  = CuArray(A)
        d_B  = CuArray(B)

        d_F  = cholesky(d_A, Val(false))
        F    = cholesky(A, Val(false))
        @test F.U   ≈ collect(d_F.U)
        @test F\(A'B) ≈ collect(d_F\(d_A'd_B))

        d_F  = cholesky(Hermitian(d_A, :L), Val(false))
        F    = cholesky(Hermitian(A, :L), Val(false))
        @test F.L   ≈ collect(d_F.L)
        @test F\(A'B) ≈ collect(d_F\(d_A'd_B))

        @test_throws DimensionMismatch LinearAlgebra.LAPACK.potrs!('U',d_A,CUDA.rand(elty,m,m))

        A    = rand(elty,m,n)
        d_A  = CuArray(A)
        @test_throws DimensionMismatch cholesky(d_A)
        @test_throws DimensionMismatch LinearAlgebra.LAPACK.potrs!('U',d_A,d_B)

        A    = zeros(elty,n,n)
        d_A  = CuArray(A)
        @test_throws LinearAlgebra.PosDefException cholesky(d_A)
    end

    CUDA.CUSOLVER.version() >= v"10.1" && @testset "Cholesky inverse (potri)" begin
        # test lower
        A    = rand(elty,n,n)
        A    = A*A' #posdef
        d_A  = CuArray(A)

        LinearAlgebra.LAPACK.potrf!('L', A)
        LinearAlgebra.LAPACK.potrf!('L', d_A)

        LinearAlgebra.LAPACK.potri!('L', A)
        LinearAlgebra.LAPACK.potri!('L', d_A)
        @test A  ≈ collect(d_A)

        # test upper
        A    = rand(elty,n,n)
        A    = A*A' #posdef
        d_A  = CuArray(A)

        LinearAlgebra.LAPACK.potrf!('U', A)
        LinearAlgebra.LAPACK.potrf!('U', d_A)
        LinearAlgebra.LAPACK.potri!('U', A)
        LinearAlgebra.LAPACK.potri!('U', d_A)
        @test A  ≈ collect(d_A)
    end

    @testset "getrf!" begin
        A          = rand(elty,m,n)
        d_A        = CuArray(A)
        d_A,d_ipiv = CUSOLVER.getrf!(d_A)
        h_A        = collect(d_A)
        h_ipiv     = collect(d_ipiv)
        alu        = LinearAlgebra.LU(h_A, convert(Vector{BlasInt},h_ipiv), zero(BlasInt))
        @test A ≈ Array(alu)
        A    = zeros(elty,n,n)
        d_A  = CuArray(A)
        @test_throws LinearAlgebra.SingularException CUSOLVER.getrf!(d_A)
    end

    @testset "getrs!" begin
        A          = rand(elty,n,n)
        d_A        = CuArray(A)
        d_A,d_ipiv = CUSOLVER.getrf!(d_A)
        B          = rand(elty,n,n)
        d_B        = CuArray(B)
        d_B        = CUSOLVER.getrs!('N',d_A,d_ipiv,d_B)
        h_B        = collect(d_B)
        @test h_B  ≈ A\B
        A          = rand(elty,m,n)
        d_A        = CuArray(A)
        @test_throws DimensionMismatch CUSOLVER.getrs!('N',d_A,d_ipiv,d_B)
        A          = rand(elty,n,n)
        d_A        = CuArray(A)
        B          = rand(elty,m,n)
        d_B        = CuArray(B)
        @test_throws DimensionMismatch CUSOLVER.getrs!('N',d_A,d_ipiv,d_B)
    end

    @testset "Matrix division" begin
        A = rand(elty,n,n)
        B = rand(elty,n,n)
        C = A \ B
        d_A = CuArray(A)
        d_B = CuArray(B)
        @test C ≈ Array(d_A \ d_B)
    end

    @testset "geqrf!" begin
        A         = rand(elty,m,n)
        d_A       = CuArray(A)
        d_A,d_tau = CUSOLVER.geqrf!(d_A)
        h_A       = collect(d_A)
        h_tau     = collect(d_tau)
        qra       = LinearAlgebra.QR(h_A, h_tau)
        @test A ≈ Array(qra)
    end

    @testset "ormqr!" begin
        A          = rand(elty, m, n)
        d_A        = CuArray(A)
        d_A, d_tau = CUSOLVER.geqrf!(d_A)
        B          = rand(elty, n, l)
        d_B        = CuArray(B)
        d_B        = CUSOLVER.ormqr!('L', 'N', d_A, d_tau, d_B)
        h_B        = collect(d_B)
        F          = qr!(A)
        @test h_B  ≈ Array(F.Q)*B
        A          = rand(elty, n, m)
        d_A        = CuArray(A)
        d_A, d_tau = CUSOLVER.geqrf!(d_A)
        B          = rand(elty, n, l)
        d_B        = CuArray(B)
        d_B        = CUSOLVER.ormqr!('L', 'N', d_A, d_tau, d_B)
        h_B        = collect(d_B)
        F          = qr!(A)
        @test h_B  ≈ Array(F.Q)*B
        A          = rand(elty, m, n)
        d_A        = CuArray(A)
        d_A, d_tau = CUSOLVER.geqrf!(d_A)
        B          = rand(elty, l, m)
        d_B        = CuArray(B)
        d_B        = CUSOLVER.ormqr!('R', 'N', d_A, d_tau, d_B)
        h_B        = collect(d_B)
        F          = qr!(A)
        @test h_B  ≈ B*Array(F.Q)
        A          = rand(elty, n, m)
        d_A        = CuArray(A)
        d_A, d_tau = CUSOLVER.geqrf!(d_A)
        B          = rand(elty, l, n)
        d_B        = CuArray(B)
        d_B        = CUSOLVER.ormqr!('R', 'N', d_A, d_tau, d_B)
        h_B        = collect(d_B)
        F          = qr!(A)
        @test h_B  ≈ B*Array(F.Q)
    end

    @testset "orgqr!" begin
        A         = rand(elty,n,m)
        d_A       = CuArray(A)
        d_A,d_tau = CUSOLVER.geqrf!(d_A)
        d_Q       = CUSOLVER.orgqr!(d_A, d_tau)
        h_Q       = collect(d_Q)
        F         = qr!(A)
        @test h_Q ≈ Array(F.Q)
        A         = rand(elty,m,n)
        d_A       = CuArray(A)
        d_A,d_tau = CUSOLVER.geqrf!(d_A)
        d_Q       = CUSOLVER.orgqr!(d_A, d_tau)
        h_Q       = collect(d_Q)
        F         = qr!(A)
        @test h_Q ≈ Array(F.Q)
    end

    @testset "sytrf!" begin
        A          = rand(elty,n,n)
        A          = A + A' #symmetric
        d_A        = CuArray(A)
        d_A,d_ipiv = CUSOLVER.sytrf!('U',d_A)
        h_A        = collect(d_A)
        h_ipiv     = collect(d_ipiv)
        A, ipiv    = LAPACK.sytrf!('U',A)
        @test ipiv == h_ipiv
        @test A ≈ h_A
        A    = rand(elty,m,n)
        d_A  = CuArray(A)
        @test_throws DimensionMismatch CUSOLVER.sytrf!('U',d_A)
        A    = zeros(elty,n,n)
        d_A  = CuArray(A)
        @test_throws LinearAlgebra.SingularException CUSOLVER.sytrf!('U',d_A)
    end

    @testset "gebrd!" begin
        A                             = rand(elty,m,n)
        d_A                           = CuArray(A)
        d_A, d_D, d_E, d_TAUQ, d_TAUP = CUSOLVER.gebrd!(d_A)
        h_A                           = collect(d_A)
        h_D                           = collect(d_D)
        h_E                           = collect(d_E)
        h_TAUQ                        = collect(d_TAUQ)
        h_TAUP                        = collect(d_TAUP)
        A,d,e,q,p                     = LAPACK.gebrd!(A)
        @test A ≈ h_A
        @test d ≈ h_D
        @test e[min(m,n)-1] ≈ h_E[min(m,n)-1]
        @test q ≈ h_TAUQ
        @test p ≈ h_TAUP
    end

    @testset "syevd!" begin
        A              = rand(elty,m,m)
        A             += A'
        d_A            = CuArray(A)
        local d_W, d_V
        if( elty <: Complex )
            d_W, d_V   = CUSOLVER.heevd!('V','U', d_A)
        else
            d_W, d_V   = CUSOLVER.syevd!('V','U', d_A)
        end
        h_W            = collect(d_W)
        h_V            = collect(d_V)
        Eig            = eigen(A)
        @test Eig.values ≈ h_W
        @test abs.(Eig.vectors'*h_V) ≈ I
        d_A            = CuArray(A)
        if( elty <: Complex )
            d_W   = CUSOLVER.heevd!('N','U', d_A)
        else
            d_W   = CUSOLVER.syevd!('N','U', d_A)
        end
        h_W            = collect(d_W)
        @test Eig.values ≈ h_W
    end

    @testset "sygvd!" begin
        A              = rand(elty,m,m)
        B              = rand(elty,m,m)
        A             *= A'
        B             *= B'
        d_A            = CuArray(A)
        d_B            = CuArray(B)
        local d_W, d_VA, d_VB
        if( elty <: Complex )
            d_W, d_VA, d_VB = CUSOLVER.hegvd!(1, 'V','U', d_A, d_B)
        else
            d_W, d_VA, d_VB = CUSOLVER.sygvd!(1, 'V','U', d_A, d_B)
        end
        h_W            = collect(d_W)
        h_VA           = collect(d_VA)
        h_VB           = collect(d_VB)
        Eig            = eigen(Hermitian(A), Hermitian(B))
        @test Eig.values ≈ h_W
        @test A*h_VA ≈ B*h_VA*Diagonal(h_W) rtol=1e-4
        # test normalization condition for eigtype 1
        @test abs.(h_VA'B*h_VA) ≈ Matrix(one(elty)*I, m, m)
        d_A            = CuArray(A)
        d_B            = CuArray(B)
        if( elty <: Complex )
            d_W   = CUSOLVER.hegvd!(1, 'N','U', d_A, d_B)
        else
            d_W   = CUSOLVER.sygvd!(1, 'N','U', d_A, d_B)
        end
        h_W            = collect(d_W)
        @test Eig.values ≈ h_W
        d_B            = CUDA.rand(elty, m+1, m+1)
        if( elty <: Complex )
            @test_throws DimensionMismatch CUSOLVER.hegvd!(1, 'N','U', d_A, d_B)
        else
            @test_throws DimensionMismatch CUSOLVER.sygvd!(1, 'N','U', d_A, d_B)
        end
    end

    @testset "syevj!" begin
        A              = rand(elty,m,m)
        B              = rand(elty,m,m)
        A             *= A'
        B             *= B'
        d_A            = CuArray(A)
        d_B            = CuArray(B)
        local d_W, d_VA, d_VB
        if( elty <: Complex )
            d_W, d_VA, d_VB = CUSOLVER.hegvj!(1, 'V','U', d_A, d_B)
        else
            d_W, d_VA, d_VB = CUSOLVER.sygvj!(1, 'V','U', d_A, d_B)
        end
        h_W            = collect(d_W)
        h_VA           = collect(d_VA)
        h_VB           = collect(d_VB)
        Eig            = eigen(Hermitian(A), Hermitian(B))
        @test Eig.values ≈ h_W
        @test A*h_VA ≈ B*h_VA*Diagonal(h_W) rtol=1e-4
        # test normalization condition for eigtype 1
        @test abs.(h_VA'B*h_VA) ≈ Matrix(one(elty)*I, m, m)
        d_A            = CuArray(A)
        d_B            = CuArray(B)
        if( elty <: Complex )
            d_W   = CUSOLVER.hegvj!(1, 'N','U', d_A, d_B)
        else
            d_W   = CUSOLVER.sygvj!(1, 'N','U', d_A, d_B)
        end
        h_W            = collect(d_W)
        @test Eig.values ≈ h_W
    end

    @testset "elty = $elty" for elty in [Float32, Float64, ComplexF32, ComplexF64]
        @testset "syevjBatched!" begin
            # Generate a random symmetric/hermitian matrix
            A = rand(elty, m,m,n)
            A += permutedims(A, (2,1,3))
            d_A = CuArray(A)

            # Run the solver
            local d_W, d_V
            if( elty <: Complex )
                d_W, d_V   = CUSOLVER.heevjBatched!('V','U', d_A)
            else
                d_W, d_V   = CUSOLVER.syevjBatched!('V','U', d_A)
            end

            # Pull it back to hardware
            h_W   = collect(d_W)
            h_V   = collect(d_V)

            # Use non-GPU blas to estimate the eigenvalues as well
            for i = 1:n
                # Get our eigenvalues
                Eig = eigen(LinearAlgebra.Hermitian(A[:,:,i]))

                # Compare to the actual ones
                @test Eig.values ≈ h_W[:,i]
                @test abs.(Eig.vectors'*h_V[:,:,i]) ≈ I
            end

            # Do it all again, but with the option to not compute eigenvectors
            d_A = CuArray(A)

            # Run the solver
            local d_W
            if( elty <: Complex )
                d_W   = CUSOLVER.heevjBatched!('N','U', d_A)
            else
                d_W   = CUSOLVER.syevjBatched!('N','U', d_A)
            end

            # Pull it back to hardware
            h_W   = collect(d_W)

            # Use non-GPU blas to estimate the eigenvalues as well
            for i = 1:n
                # Get the reference results
                Eig = eigen(LinearAlgebra.Hermitian(A[:,:,i]))

                # Compare to the actual ones
                @test Eig.values ≈ h_W[:,i]
            end
        end
    end

    @testset "svd with $method method" for
        method in (CUSOLVER.QRAlgorithm, CUSOLVER.JacobiAlgorithm),
        (_m, _n) in ((m, n), (n, m))

        A              = rand(elty, _m, _n)
        U, S, V        = svd(A, full=true)
        d_A            = CuArray(A)

        if _m > _n || method == CUSOLVER.JacobiAlgorithm
            d_U, d_S, d_V  = svd(d_A, method, full=true)
            h_S            = collect(d_S)
            h_U            = collect(d_U)
            h_V            = collect(d_V)
            @test abs.(h_U'h_U) ≈ I
            @test abs.(h_U[:,1:min(_m,_n)]'U[:,1:min(_m,_n)]) ≈ I
            @test collect(svdvals(d_A, method)) ≈ svdvals(A)
            @test abs.(h_V'*h_V) ≈ I
            @test abs.(h_V[:,1:min(_m,_n)]'*V[:,1:min(_m,_n)]) ≈ I
            @test collect(d_U'*d_A*d_V) ≈ U'*A*V
            @test collect(svd(d_A, method).V') == h_V[:,1:min(_m,_n)]'
        else
            @test_throws ArgumentError svd(d_A, method)
        end
    end
    # Check that constant propagation works
    _svd(A) = svd(A, CUSOLVER.QRAlgorithm)
    @inferred _svd(CUDA.rand(Float32, 4, 4))

    @testset "qr" begin
        tol = min(m, n)*eps(real(elty))*(1 + (elty <: Complex))

        A              = rand(elty, m, n)
        qra            = qr(A)
        d_A            = CuArray(A)
        d_F            = qr(d_A)
        d_RR           = d_F.Q'*d_A
        @test d_RR[1:n,:] ≈ d_F.R atol=tol*norm(A)
        @test norm(d_RR[n+1:end,:]) < tol*norm(A)
        @test size(d_F) == size(A)
        @test size(d_F.Q, 1) == size(A, 1)
        @test det(d_F.Q) ≈ det(collect(d_F.Q * CuMatrix{elty}(I, size(d_F.Q)))) atol=tol*norm(A)
        CUDA.@allowscalar begin
            qval = d_F.Q[1, 1]
            @test qval ≈ qra.Q[1, 1]
            qrstr = sprint(show, d_F)
            @test qrstr == "$(typeof(d_F)) with factors Q and R:\n$(sprint(show, d_F.Q))\n$(sprint(show, d_F.R))"
        end
        dQ, dR = d_F
        @test collect(dQ*dR) ≈ A
        A              = rand(elty, n, m)
        d_A            = CuArray(A)
        d_F            = qr(d_A)
        @test d_F.Q'*d_A ≈ d_F.R atol=tol*norm(A)
        @test det(d_F.Q) ≈ det(collect(d_F.Q * CuMatrix{elty}(I, size(d_F.Q)))) atol=tol*norm(A)
        A              = rand(elty, m, n)
        d_A            = CuArray(A)
        h_q, h_r       = qr(d_A)
        q, r           = qr(A)
        @test Array(h_q) ≈ Array(q)
        @test collect(CuArray(h_q)) ≈ Array(q)
        @test Array(h_r) ≈ Array(r)
        A              = rand(elty, n, m)
        d_A            = CuArray(A)
        h_q, h_r       = qr(d_A) # FixMe! Use iteration protocol when implemented
        q, r           = qr(A)
        @test Array(h_q) ≈ Array(q)
        @test Array(h_r) ≈ Array(r)
    end

    @testset "potrsBatched!" begin
        @testset "elty = $elty" for elty in [Float32, Float64, ComplexF32, ComplexF64]
            # Test lower
            bA = [rand(elty, m, m) for i in 1:n]
            bA = [bA[i]*bA[i]' for i in 1:n]
            bB = [rand(elty, m) for i in 1:n]

            # move to device
            bd_A = CuArray{elty, 2}[]
            bd_B = CuArray{elty, 1}[]
            for i in 1:length(bA)
                push!(bd_A, CuArray(bA[i]))
                push!(bd_B, CuArray(bB[i]))
            end

            bd_X = CUSOLVER.potrsBatched!('L', bd_A, bd_B)
            bh_X = [collect(bd_X[i]) for i in 1:n]

            for i = 1:n
                LinearAlgebra.LAPACK.potrs!('L', bA[i], bB[i])
                @test bB[i] ≈ bh_X[i]
            end

            # Test upper
            bA = [rand(elty, m, m) for i in 1:n]
            bA = [bA[i]*bA[i]' for i in 1:n]
            bB = [rand(elty, m) for i in 1:n]

            # move to device
            bd_A = CuArray{elty, 2}[]
            bd_B = CuArray{elty, 1}[]
            for i in 1:length(bA)
                push!(bd_A, CuArray(bA[i]))
                push!(bd_B, CuArray(bB[i]))
            end

            bd_X = CUSOLVER.potrsBatched!('U', bd_A, bd_B)
            bh_X = [collect(bd_X[i]) for i in 1:n]

            for i = 1:n
                LinearAlgebra.LAPACK.potrs!('U', bA[i], bB[i])
                @test bB[i] ≈ bh_X[i]
            end
            # error throwing tests
            bA = [rand(elty, m, m) for i in 1:n]
            bA = [bA[i]*bA[i]' for i in 1:n]
            bB = [rand(elty, m) for i in 1:n+1]

            # move to device
            bd_A = CuArray{elty, 2}[]
            bd_B = CuArray{elty, 1}[]
            for i in 1:length(bA)
                push!(bd_A, CuArray(bA[i]))
                push!(bd_B, CuArray(bB[i]))
            end
            push!(bd_B, CuArray(bB[end]))

            @test_throws DimensionMismatch CUSOLVER.potrsBatched!('L', bd_A, bd_B)

            bA = [rand(elty, m, m) for i in 1:n]
            bA = [bA[i]*bA[i]' for i in 1:n]
            bB = [rand(elty, m) for i in 1:n]
            bB[1] = rand(elty, m+1)
            # move to device
            bd_A = CuArray{elty, 2}[]
            bd_B = CuArray{elty, 1}[]
            for i in 1:length(bA)
                push!(bd_A, CuArray(bA[i]))
                push!(bd_B, CuArray(bB[i]))
            end

            @test_throws DimensionMismatch CUSOLVER.potrsBatched!('L', bd_A, bd_B)

            bA = [rand(elty, m, m) for i in 1:n]
            bA = [bA[i]*bA[i]' for i in 1:n]
            bB = [rand(elty, m, m) for i in 1:n]
            # move to device
            bd_A = CuArray{elty, 2}[]
            bd_B = CuArray{elty, 2}[]
            for i in 1:length(bA)
                push!(bd_A, CuArray(bA[i]))
                push!(bd_B, CuArray(bB[i]))
            end

            @test_throws ArgumentError CUSOLVER.potrsBatched!('L', bd_A, bd_B)
        end
    end

    @testset "potrfBatched!" begin
        @testset "elty = $elty" for elty in [Float32, Float64, ComplexF32, ComplexF64]
            # Test lower
            bA = [rand(elty, m, m) for i in 1:n]
            bA = [bA[i]*bA[i]' for i in 1:n]

            # move to device
            bd_A = CuArray{elty, 2}[]
            for i in 1:length(bA)
                push!(bd_A, CuArray(bA[i]))
            end

            bd_A, info = CUSOLVER.potrfBatched!('L', bd_A)
            bh_A = [collect(bd_A[i]) for i in 1:n]

            for i = 1:n
                LinearAlgebra.LAPACK.potrf!('L', bA[i])
                @test bA[i] ≈ bh_A[i]
            end

            # Test upper
            bA = [rand(elty, m, m) for i in 1:n]
            bA = [bA[i]*bA[i]' for i in 1:n]

            # move to device
            bd_A = CuArray{elty, 2}[]
            for i in 1:length(bA)
                push!(bd_A, CuArray(bA[i]))
            end

            bd_A, info = CUSOLVER.potrfBatched!('U', bd_A)
            bh_A = [collect(bd_A[i]) for i in 1:n]

            for i = 1:n
                LinearAlgebra.LAPACK.potrf!('U', bA[i])
                # cuSOLVER seems to return symmetric/hermitian matrix when using 'U'
                @test Hermitian(bA[i]) ≈ bh_A[i]
            end
        end
    end

end


using Meshing
using GeometryTypes
using LinearAlgebra: dot, norm
using FileIO

f(v) = sqrt(sum(dot(v,v))) - 1
sdf = SignedDistanceField(f,HyperRectangle(Vec(-1,-1,-1.), Vec(2,2,2.)))

mc = HomogenousMesh(sdf, MarchingCubes())
mt = HomogenousMesh(sdf, MarchingTetrahedra())
ns = HomogenousMesh(sdf, NaiveSurfaceNets())

# save the Sphere as a PLY file
save("sphere_mc.ply",mc)
save("sphere_mt.ply",mt)
save("sphere_ns.ply",ns)


"""
    Mountain Car

    https://en.wikipedia.org/wiki/Mountain_car_problem
"""

struct MountainCar{I, T} <: Model{I, T}
    n::Int
    m::Int
    d::Int

    gravity
    mass
    steepness
    friction

    xl
    xu
    ul
    uu
end

function fd(model::MountainCar{Discrete, FixedTime}, x⁺, x, u, w, h, t)
    x⁺ - [x[1] + x[2] * h;
          x[2] + (model.gravity * model.mass * cos(model.steepness * x[1])
            + u[1] / model.mass - model.friction * x[2]) * h]
end

function fd(model::MountainCar{Discrete, FixedTime}, x, u, w, h, t)
    [x[1] + x[2] * h;
     x[2] + (model.gravity * model.mass * cos(model.steepness * x[1])
        + u[1] / model.mass - model.friction * x[2]) * h]
end

n, m, d = 2, 1, 0
gravity = 9.81
mass = 0.2
steepness = 3.0
friction = 0.3
xl = [-1.2; -1.5]
xu = [0.5; 1.5]
ul = [-2.0]
uu = [2.0]
model = MountainCar{Discrete, FixedTime}(n, m, d,
    gravity, mass, steepness, friction, xl, xu, ul, uu)

function visualize!(vis, model::Car, x;
       Δt = 0.1, traj = false)

	default_background!(vis)

	obj_path = joinpath(pwd(),"models/cybertruck/cybertruck.obj")
	mtl_path = joinpath(pwd(),"models/cybertruck/cybertruck.mtl")

	m = ModifiedMeshFileObject(obj_path,mtl_path,scale=0.05)
	setobject!(vis["car"],m)

   	anim = MeshCat.Animation(convert(Int, floor(1.0 / Δt)))

    for t = 1:length(x)
        MeshCat.atframe(anim,t) do
			settransform!(vis["car"],
				compose(Translation(x[t][1], x[t][2], 0.0),
					LinearMap(RotZ(x[t][3] + pi) * RotX(pi / 2.0))))
        end
    end

	if traj
		pts_nom = collect(eachcol(hcat(x...)))
		material_nom = LineBasicMaterial(color = colorant"orange", linewidth = 3.0)
		setobject!(vis["traj"], Object(PointCloud(pts_nom), material_nom, "Line"))
	end

	settransform!(vis["/Cameras/default"],
		compose(Translation(0.0, 0.0, 0.0),LinearMap(RotY(-pi/2.5))))

    MeshCat.setanimation!(vis, anim)
end


# Multivariate Normal distribution

import StatsFuns
using LinearAlgebra
export MvNormal
using Random
import Base
using Tullio

struct MvNormal{N, T, I, J} <: ParameterizedMeasure{N}
    par::NamedTuple{N, T}
end

function MvNormal(nt::NamedTuple{N,T}) where {N,T}
    I,J = mvnormaldims(nt)

    cache = Vector{Float64}(undef, max(I,J))
    MvNormal{N,T,I,J}(cache, nt)
end

function Base.size(d::MvNormal{N, T, I, J}) where {N,T,I,J}
    return (I,)
end

mvnormaldims(nt::NamedTuple{(:A, :b)}) = size(nt.A)

function MeasureTheory.basemeasure(μ::MvNormal{N, T, I,I}) where {N, T, I}
    return (1 / sqrt2π)^I * Lebesgue(ℝ)^I
end

sampletype(d::MvNormal{N, T, I, J}) where {N,T,I,J} = Vector{Float64}

MvNormal(; kwargs...) = begin
    MvNormal((; kwargs...))
end

function Random.rand!(rng::AbstractRNG, d::MvNormal{(:A, :b),T,I,J}, x::AbstractArray) where {T,I,J}
    z = getcache(d, J)
    rand!(rng, Normal()^J, z)
    copyto!(x, d.b)
    mul!(x, d.A, z, 1.0, 1.0)
    return x
end

function logdensity(d::MvNormal{(:A,:b)}, x)
    cache = getcache(d, size(d))
    z = d.A \ (x - d.b)
    return logdensity(MvNormal(), z) - logabsdet(d.A)
end

≪(::MvNormal, ::Lebesgue{ℝ}) = true

function logdensity(d::MvNormal{(:Σ⁻¹,)}, x)
    @tullio ℓ = -0.5 * x[i] * d.Σ⁻¹[i,j] * x[j]
    return ℓ
end

mvnormaldims(nt::NamedTuple{(:Σ⁻¹,)}) = size(nt.Σ⁻¹)


using Combinatorics

"""
one_body_to_matrix and two_body_to_matrix for testing purposes
returns the 2^L×2^L matrices of a_k^†a_l and a_k^† a_l^† a_m a_n
"""
function one_body_to_matrix(k,l,L)
    A = zeros(Float64,2^L,2^L)
    for i in 0:2^L-1
        i_bin = digits(i,base=2,pad=L)
        for j in 0:2^L-1
            j_bin = digits(j,base=2,pad=L) #corresponding index as d-tuples
            if i_bin[k]==1 && j_bin[l]==1 && j-2^(l-1) == i-2^(k-1)
                A[i+1,j+1] = (-1)^(sum(i_bin[1:(k-1)])+sum(j_bin[1:(l-1)]))
            end
        end
    end
    return A
end

function two_body_to_matrix(k,l,m,n,L)
    A = zeros(Float64,2^L,2^L)
    for i in 0:2^L-1
        i_bin = digits(i,base=2,pad=L)
        for j in 0:2^L-1
            j_bin = digits(j,base=2,pad=L) #corresponding index as d-tuples
            if i_bin[k]==1 && i_bin[l]==1 && j_bin[m]==1 && j_bin[n]==1 && j-2^(m-1)-2^(n-1) == i-2^(k-1)-2^(l-1)
                A[i+1,j+1] = (-1)^(sum(i_bin[1:(k-1)])+sum(i_bin[1:(l-1)])+sum(j_bin[1:(m-1)])+sum(j_bin[1:(n-1)]))
            end
        end
    end
    return A
end

#naive implementation

"""
h: two-dimensional array of the connected sites (kinetic+external potential)
V: four-dimensional array of the connected sites (interaction)
"""
function sc_to_mat(h,V)
    L = size(h,1)
    A = zeros(2^L,2^L)
    for i in 1:L
        for j in i:L
            for k in j:L
                for l in k:L
                    if !isapprox(h[i,j],0.0,atol=1e-10)
                        H = one_body_to_matrix(i,j,L)
                        A += h[i,j]*(H+H')
                    end
                    if !isapprox(V[i,j,k,l],0.0,atol=1e-10)
                        H = two_body_to_matrix(i,j,k,l,L)
                        A += V[i,j,k,l]*(H+H')
                    end
                end
            end
        end
    end
    return A
end

#odd index = spin up, even index = spin down
#returns TT of a half-filled state
function half_filling(N)
    tt_vec = Array{Array{Float64,3},1}(undef,2N)
    for i in 1:N
        tt_vec[2i-1] = zeros(2,1,1)
        tt_vec[2i-1][1,1,1] = 1.0
        tt_vec[2i] = zeros(2,1,1)
        tt_vec[2i][2,1,1] = 1.0
    end
    return ttvector{Float64}(tt_vec,2*ones(Int,2N),ones(Int,2N+1),zeros(2N))
end

#auxiliary functions for a_p^†a_q
function mpo_core_id(;T=Float64)
    out = zeros(T,2,2,1,1)
    out[1,1,1,1] = 1.0
    out[2,2,1,1] = 1.0
    return out
end

function mpo_core_ferm_sign(;T=Float64)
    out = zeros(T,2,2,1,1)
    out[1,1,1,1] = 1.0
    out[2,2,1,1] = -1.0
    return out
end

function mpo_core_creation(;T=Float64)
    out = zeros(T,2,2,1,1)
    out[2,1,1,1] = 1.0
    return out
end

function mpo_core_annihilation(;T=Float64)
    out = zeros(T,2,2,1,1)
    out[1,2,1,1] = 1.0
    return out
end

#returns bosonic or fermionic MPO of a_p^†a_q
function one_body_mpo(p,q,L;fermion=true,T=Float64)
    H = Array{Array{T,4},1}(undef,L)
    if p == q
        for i in 1:L
            H[i] = mpo_core_id(;T=T)
        end
        H[p][1,1,1,1] = 0.0
    else
        H[p] = mpo_core_creation(;T=T)
        H[q] = mpo_core_annihilation(;T=T)
        for i in 1:min(p,q)-1
            H[i] = mpo_core_id(;T=T)
        end
        for i in max(p,q)+1:L
            H[i] = mpo_core_id(;T=T)
        end
        for i in min(p,q)+1:max(p,q)-1
            if fermion
                H[i] = mpo_core_ferm_sign(;T=T)
            else
                H[i] = mpo_core_id(;T=T)
            end
        end
    end
    return ttoperator{T}(H,2*ones(Int64,L),ones(Int64,L+1),zeros(Int64,L))
end

#returns bosonic or fermionic MPO of a_k^† a^†_l a_m a_n
#assume k<l,m<n
function two_body_mpo(k,l,m,n,L;T=Float64)
    H = Array{Array{T,4},1}(undef,L)
    if l == m
        A = one_body_mpo(k,m,L;T=T)
        B = one_body_mpo(l,n,L;T=T)
        C = one_body_mpo(k,n,L;T=T)
        return  (-1.0*A)*B + C
    else
        A = one_body_mpo(k,m,L;T=T)
        B = one_body_mpo(l,n,L;T=T)
        return -1.0*(A*B)
    end
end

"""
returns an MPO version of
H = Σ_{ij} h_ij (a_i† a_j + c.c.) + Σ_{ijkl} V_{ijkl} (a_i†a_j†a_ka_l + c.c.)
"""
#assuming diagonal terms are divided by 2 in the h and V matrix
function hV_to_mpo(h,V;tol=1e-8)
    L = size(h,1)
    i_list = findall(!iszero,h)
    if length(i_list)>0
        i = i_list[1]
        A = h[i]*(one_body_mpo(i[1],i[2],L)+one_body_mpo(i[2],i[1],L))
        for i in i_list[2:end]
            H = one_body_mpo(i[1],i[2],L)
            G = one_body_mpo(i[2],i[1],L)
            A = A+h[i]*(G+H)
        end
        A = tt_rounding(A,tol=tol)
    end
    for i in findall(!iszero,V)
        H = two_body_mpo(i[1],i[2],i[3],i[4],L)
        G = two_body_mpo(i[4],i[3],i[2],i[1],L)
        A = tt_rounding(A+V[i]*(H+G),tol=tol)
    end
    return A
end

"""
Examples of standard Hamiltonians
"""

#returns the h,V for the hubbard model (assuming natural labelling of sites)
function hubbard_1D(L;t=1,U=1, pbc=false)
    h = zeros(2L,2L)
    V = zeros(2L,2L,2L,2L)
    for i in 1:L-1
        h[2i,2i+2]=-t
        h[2i-1,2i+1]=-t
        V[2i-1,2i,2i-1,2i]=-U #because of the anticommutation rules
    end
    V[2L-1,2L,2L-1,2L]=-U
    if pbc
        h[2L-1,1]=-t
        h[2L,2]=-t
    end
    return h,V
end

#switch sites j and k in a one-dimensional spin chain model
function site_switch(j::Integer,k::Integer,L::Integer)
    out = collect(1:2L)
    out[2j-1],out[2j] = 2k-1,2k
    out[2k-1],out[2k] = 2j-1,2j
    return out
end

#2D Hubbard with cylindrical boundary conditions along L
function hubbard_2D(w,L;t=1,U=1,w_pbc = false,L_pbc=true)
    h = zeros(2w*L,2w*L)
    V = zeros(2w*L,2w*L,2w*L,2w*L)
    for i in 0:w-2
        for j in 0:L-2
            V[2(j*w+i)+1,2(j*w+i)+2,2(j*w+i)+1,2(j*w+i)+2] = -U
            h[2(j*w+i)+1,2((j+1)*w+i)+1] = -t #right connection
            h[2(j*w+i)+1,2(j*w+i+1)+1] = -t #bottom connection
            h[2(j*w+i)+2,2((j+1)*w+i)+2] = -t #right connection
            h[2(j*w+i)+2,2(j*w+i+1)+2] = -t #bottom connection
        end
        #j=L-1
        V[2((L-1)*w+i)+1,2((L-1)*w+i)+2,2((L-1)*w+i)+1,2((L-1)*w+i)+2] = -U
        h[2((L-1)*w+i)+1,2((L-1)*w+i+1)+1] = -t #bottom connection
        h[2((L-1)*w+i)+2,2((L-1)*w+i+1)+2] = -t #bottom connection
        if L_pbc && L>2
            h[2(L*w+i-w)+1,2i+1] =-t #right connection for j=L-1
            h[2(L*w+i-w)+2,2i+2] =-t #right connection for j=L-1
        end
    end
    #i = w-1
    for j in 0:L-2
        V[2(j*w+w-1)+1,2(j*w+w-1)+2,2(j*w+w-1)+1,2(j*w+w-1)+2] = -U
        h[2(j*w+w-1)+1,2((j+1)*w+w-1)+1] = -t #right connection
        h[2(j*w+w-1)+2,2((j+1)*w+w-1)+2] = -t #right connection
        if w_pbc && w>2
            h[2(j*w+w-1)+1,2j*w+1] = -t #top connection
            h[2(j*w+w-1)+2,2j*w+2] = -t #top connection
        end
    end
    #i = w-1, j=L-1
    V[2L*w-1,2L*w,2L*w-1,2L*w] = -U
    if L_pbc && L>2
        h[2(L*w-1)+1,2(w-1)+1] = -t #right connection
        h[2(L*w-1)+2,2(w-1)+2] = -t #right connection
    end
    if w_pbc && w>2
        h[2(L*w-1)+1,2(L-1)*w+1] = -t #right connection
        h[2(L*w-1)+2,2(L-1)*w+2] = -t #right connection
    end
    return h,V
end

"""
PPP Hamiltonian for cyclic polyene C_NH_N (ref. G. Fano, F. Ortolani † and L. Ziosi, The density matrix renormalization group method. Application to the PPP model of a cyclic polyene chain, J. Chem. Phys. 1998)
H = β ∑_{<μ,ν>,σ} a^†_μσ a^†_νσ + c.c. + 0.5 ∑_μ,ν γ_μν (n_μ-1)(n_ν-1)

β = -2.5eV, γ_μν = 1/(γ0^-1+d_μν), γ0 = 10.84eV
d_μν = b sin(π/N*(μ-ν [N]))/sin(π/N), b = 1.4 A

Ground-state energy = -12.72 eV (N=6)
"""
function PPP_C_NH_N(N;β=-2.5/27.2113845,b=1.4*1.8897259886,γ0=10.84/27.2113845,order=collect(1:2N),tol=1e-8)
    @assert isperm(order) "Ordering given is not a permutation."
    h = zeros(2N,2N)
    γ = sum(1/(1/γ0+b*sin(k/N*pi)/sin(pi/N)) for k in 1:N)
    for i in 1:N-1
        h[2i-1,2i+1] = β
        h[2i,2i+2] = β
    end
    h[2N,2] = β
    h[2N-1,1] = β
    H_tto = hV_to_mpo(h[order,order],zeros(2N,2N,2N,2N))
    σ = invperm(order)
    id = -1.0*id_tto(2N)
    for i in 1:N
        Hi = (one_body_mpo(σ[2i],σ[2i],2N)+one_body_mpo(σ[2i-1],σ[2i-1],2N))+id
        for j in 1:N
            Hj = (one_body_mpo(σ[2j],σ[2j],2N)+one_body_mpo(σ[2j-1],σ[2j-1],2N)) + id
            Htemp = Hi*Hj
            γij = 1/(1/γ0+b*sin(pi/N*mod(i-j,N))/sin(pi/N))
            H_tto = H_tto + 0.5*γij*Htemp
            H_tto = tt_rounding(H_tto,tol=tol)
        end
    end
    return H_tto
end


using Smg2s
using Test
using SparseArrays
using Plots

using LinearAlgebra
using Random

n = 10000
diag_l = -9900
diag_u = -800
nbOne = 70
nilpdiag = 3500
sp=0.0002
figname = "fig/pattern12"

nilpvec = zeros(Int64, n-1)

cnt = 0

while cnt < n-1
    l₁ = rand(1:nbOne)
    l₀ = rand(1:nbOne)

    for i = cnt+1:l₁+cnt
        if i <= n-1
            nilpvec[i] = 1
        end
    end

    global cnt += l₀ + l₁

end

#nilpMatrix = sparse(diagm(nilpdiag=>nilpvec[1:n-nilpdiag]))

#nilp = Nilp(nilpMatrix, n)
nilp = Nilp(nilpvec, nilpdiag, n)

function f(idx::Integer, n = n)
    return 10 * cos(((idx-1) / 2) * 2 * π / n) + 5
end

spec = zeros(ComplexF64, n)
Spectrum!(spec, f, n)

Am = spzeros(Float64, n, n)
initMat!(Am, diag_l, diag_u, n; scale=0.1, sparsity=sp)
@time genMat = nonsym(n, diag_l, diag_u, spec, Am, nilp)
sparsity = nnz(genMat) / (n * n)
@info "sparsity=" sparsity
spy(genMat, legend = :none, title="size:$n, diag_l:$diag_l, diag_u:$diag_u, nbOne:$nbOne, nilpdiag: $nilpdiag, sp: $sp",titlefont = font(8))


# Standard rectlin layer: f(A,x,b)=rectlin(A*x+b)
# Note that the bias (a column vector) gets expanded here to match the dimension of A*x

FabsAXplusBias(A,X,b)=begin; a=A*X+b*ones(1,size(X,2)); return (abs(a),sign(a)); end

function FabsAXplusBias_inplace(value,aux,A,X,b)
    a=A*X+b*ones(1,size(X,2))
    copy!(value,abs(a))
    copy!(aux,sign(a))
end

function DabsAXplusBias(derivativeIDX,f_c,faux_c,grad_c,grad_n,A,X,b)
    if derivativeIDX==1
        axpy!(1.0,(grad_c.*faux_c)*X',grad_n)
    elseif derivativeIDX==2
        axpy!(1.0,A'*(grad_c.*faux_c),grad_n)
    elseif derivativeIDX==3
        axpy!(1.0,sum(grad_c.*faux_c,2),grad_n)
    end
end

if PROC=="GPU"

    function FabsAXplusBias_inplace(value,aux,A::CudaArray,X::CudaArray,b::CudaArray)
        FAXplusBias_inplace(value,aux,A,X,b)
        copy!(aux,value)
        abs!(value,value)
    end

    function DabsAXplusBias(derivativeIDX,f_c,faux_c,grad_c,grad_n,A::CudaArray,X::CudaArray,b::CudaArray)
        tmp=CudaArray(Float64,size(grad_c)); fill!(tmp,0.0)
        xsigny_update!(grad_c,faux_c,tmp)
        if derivativeIDX==1
            gemm!('N','T',1.0,tmp,X,1.0,grad_n)
        elseif derivativeIDX==2
            gemm!('T','N',1.0,A,tmp,1.0,grad_n)
        elseif derivativeIDX==3
            ons=CudaArray(Float64,(size(X,2),1)); fill!(ons,1.0)
            gemm!('N','N',1.0,tmp,ons,1.0,grad_n)
            free(ons)
        end
        free(tmp)
    end

end

Derivative[FabsAXplusBias]=DabsAXplusBias
Inplace[FabsAXplusBias]=FabsAXplusBias_inplace

absAXplusBias(A,X,b)=ADnode(FabsAXplusBias,[A X b])
export absAXplusBias


#  Copyright 2017, Iain Dunning, Joey Huchette, Miles Lubin, and contributors
#  This Source Code Form is subject to the terms of the Mozilla Public
#  License, v. 2.0. If a copy of the MPL was not distributed with this
#  file, You can obtain one at http://mozilla.org/MPL/2.0/.
using JuMP, Ipopt

# Use nonlinear optimization to compute the maximum likelihood estimate (MLE)
# of the parameters of a normal distribution
# aka the sample mean and variance

n = 1000
data = randn(n)

m = Model(solver=IpoptSolver(print_level=0))

@variable(m, μ, start = 0.0)
@variable(m, σ >= 0.0, start = 1.0)

@NLobjective(m, Max, (n/2)*log(1/(2π*σ^2))-sum((data[i]-μ)^2 for i=1:n)/(2σ^2))

solve(m)

println("μ = ", getvalue(μ))
println("mean(data) = ", mean(data))
println("σ^2 = ", getvalue(σ)^2)
println("var(data) = ", var(data))
println("MLE objective: ", getobjectivevalue(m))

# constrained MLE?
@NLconstraint(m, μ == σ^2)

solve(m)
println("\nWith constraint μ == σ^2:")
println("μ = ", getvalue(μ))
println("σ^2 = ", getvalue(σ)^2)

println("Constrained MLE objective: ", getobjectivevalue(m))


abstract type AbstractConvolutionKernel end

immutable GaussianKernel <: AbstractConvolutionKernel
    Σ::AbstractPDMat
    SquaredExponentialKernel(Σ::AbstractPDMat) = new(Σ)
    SquaredExponentialKernel(Σ::Vector) = new(PDiagMat(Σ))
    SquaredExponentialKernel(Σ::Array) = new(PDMat(Σ))
end

const SquaredExponentialKernel = GaussianKernel

dim(kern::AbstractConvolutionKernel) = kern.Σ.dim

## Calculate weights
function conv_wt(kern::SquaredExponentialKernel, d::Array)
    if size(d, 1) != dim(kern)
        throw(DimensionMismatch)
    end
    exp.(-0.5 * invquad(kern.Σ, d))
end


function abseigs1d(abso::Vector{Float64},eigs::Vector{Float64},J::Matrix{Float64},
    Q::Float64,A::Float64,beta::Float64,rho::Float64,ts::Float64,xs::Float64,
    Qs::Float64,As::Float64)
    J1d(J,Q,A,beta,rho,ts,xs,Qs,As);
    eigs[1] = 0.5*(J[1,1] + J[2,2]) + sqrt.(0.25*(J[1,1] + J[2,2])^2 -
        J[1,1]*J[2,2] + J[1,2]*J[2,1]);
    eigs[2] = 0.5*(J[1,1] + J[2,2]) - sqrt.(0.25*(J[1,1] + J[2,2])^2 -
        J[1,1]*J[2,2] + J[1,2]*J[2,1]);
    abso[1] = abs(eigs[1]);
    abso[2] = abs(eigs[2]);
end


import Test: @test, @testset, @test_throws
import Luna: Capillary, Tools, PhysData

@testset "Params" begin
m = Capillary.MarcatiliMode(125e-6, :He, 0.4, model=:reduced)
p = Tools.params(300e-6, 10e-15, 800e-9, m, :He, P=0.4)
# compare to Pufe
@test isapprox(p.N, 2.239, rtol=1e-3)
@test isapprox(p.Lfiss, 1.768, rtol=1e-3)
@test isapprox(p.zdw, 378.8e-9, rtol=1e-2)
@test isapprox(p.P0/p.Pcr, 0.0398, rtol=2e-2)
# test backup zdw method
p = Tools.capillary_params(6e-9, 20e-15, 800e-9, 14e-6, :Kr, P=15.0)
@test isapprox(p.zdw, 7.693023014958748e-7, rtol=1e-7)
end

@testset "RDW phasematching" begin
a = 125e-6
gas = :HeJ
pressure = 1
λ0 = 800e-9

λRDW = Tools.λRDW(a, gas, pressure, λ0)
@test λRDW - 188e-9 < 1e-9
@test Tools.pressureRDW(a, gas, λRDW, λ0) ≈ pressure
end


function resample(n_samples, lbs, ubs, prev_X, x_sol; mean_shift=0.75, std_shrink = 0.3)
    """ Gaussian resamples free variables depending on previous solution.
    Arguments:
    - n_samples: number of samples
    - lbs: lower bounds
    - bs: upper bounds
    - prev_X: previous data
    - x_sol: solution array with same # of columns as lbs/ubs
    - sigma: by how much to move the mean to new optimum, and to shrink the std.
    Returns:
    - new_samples
    """
    X = deepcopy(prev_X);
    prev_samples, n_features = size(prev_X);
    prev_mean = [sum(X[:,i])/prev_samples for i=1:n_features];
    prev_std = [sum((X[:,i]-ones(prev_samples)* prev_mean[i]).^2)/prev_samples for i=1:n_features];
    new_std = prev_std*(1-std_shrink);
    new_mean = mean_shift*x_sol + (1-mean_shift)*prev_mean;
    dist = MvNormal(new_mean, new_std);
    while size(X,1) < prev_samples + n_samples
        new_row = rand(dist)';
        if all(lbs .<= new_row) && all(ubs .>= new_row);
            X = vcat(X, Matrix(new_row));
        end
    end
    return X
end

# Sampling based on GaussianProcesses
#         plan = randomLHC(Int(round(get_param(bbf, :n_samples)*ratio)), n_dims);
#         random_samples = scaleLHC(plan,[(bbf.lbs[i], bbf.ubs[i]) for i in vks]);
#         μ, σ = predict_f(bbf.gp, random_samples');
#         cdf_0 = [Distributions.cdf(Distributions.Normal(μ[i], σ[i]),0) for i=1:size(random_samples,1)];
#          #TODO: add criterion for information as well (something like sortperm(σ))
#         # Sample places with high probability of being near boundary (0),
#         # but also balance feasibility ratio.
#         p = bbf.feas_ratio
# #         balance_fn = x -> -1*tan(2*atan(-0.5)*(x-0.5)) + 0.5
#         balance_fn = x -> -1/0.5^2*(x-0.5)^3 + 0.5;
#         indices = sortperm(abs.(cdf_0 .- balance_fn(p)));
#         samples = DataFrame(random_samples[indices[1:get_param(bbf, :n_samples)],:], vks);
#         eval!(bbf, samples);

# function predict(bbf::BlackBoxFunction, X::AbstractArray)
#     μ, σ = predict_f(bbf.gp, transpose(X))
#     return μ, σ
# end
#
# function optimize_gp!(bbf::BlackBoxFunction)
#     """ Optimizes a GaussianProcess over a BlackBoxFunction,
#     with adaptively changing kernel. """
# #         bbf.gp = ElasticGPE(length(bbf.idxs), # data
# #         mean = MeanConst(sum(bbf.Y)/length(bbf.Y)), logNoise = -10)
#     lbs = [bbf.lbs[key] for key in bbf.vks];
#     ubs = [bbf.ubs[key] for key in bbf.vks];
#     bbf.gp = GPE(transpose(Array(bbf.X)), bbf.Y, # data
#     MeanConst(sum(bbf.Y)/length(bbf.Y)),
#     SEArd(log.((ubs-lbs)./(2*sqrt(length(bbf.Y)))), -5.))
#     optimize!(bbf.gp); #TODO: optimize GP
#                        # Instead of regenerating at every run, figure out
#                        # how to update.
# end


module NCP_VB

include("../Misc.jl")
using .Misc
import Base.Iterators: product

using Distributions, SpecialFunctions
import Base.Iterators: product

export standard_VB

function standard_VB(X::Array{ℜ,N}, R::Ƶ; a=1.0, b=a/nansum(X), EPOCHS::Ƶ=1, ϵ=1e-16) where {ℜ<:Real, Ƶ<:Integer, N}
    I = size(X)
    α_I = (a/R) ./ I
    α_R = a/R

    S_R, S₊ = zeros(R), sum(X)
    S_I = [zeros(R,Iₙ) for Iₙ in I]

    log_λ = digamma(S₊ + a) - log(b + 1)
    log_θ_R = log.(rand(Dirichlet(R,α_R + 0.1)))
    log_θ_I = [log.(rand(Dirichlet(Iₙ,α_Iₙ + 0.1),R)) for (Iₙ,α_Iₙ) in zip(I,α_I)]

    log_C = a*log(b) - (a + S₊)*log(b + 1) + lgamma(a + S₊) - lgamma(a) - sum(lgamma, X .+ 1)

    s, log_p = zeros(R), zeros(R)
    ELBO = fill(log_C,EPOCHS)

    for eph=1:EPOCHS
        # initialization
        for n ∈ 1:N
            S_I[n] .= 0
        end
        S_R .= 0

        # expectation of S
        for i ∈ product(map(Iₙ -> 1:Iₙ, I)...)
            log_p .= log_θ_R + sum(n -> log_θ_I[n][i[n],:],1:N)
            log_p .-= logsumexp(log_p)
            s .= X[i...] .* exp.(log_p)

            for (n,iₙ) ∈ enumerate(i)
                S_I[n][:,iₙ] .+= s
            end
            S_R .+= s

            ELBO[eph] -= sum(s .* log_p)
        end

        for (n,Iₙ) ∈ enumerate(I)
            log_θ_I[n] .= digamma.(S_I[n] .+ α_I[n])' .- digamma.(S_R .+ α_R)'
            ELBO[eph] += sum(lbeta(α_I[n] .+ S_I[n];dims=2)) - R*lbeta(α_I[n],Iₙ)
        end
        log_θ_R .= digamma.(S_R .+ α_R) .- digamma(S₊ + R*α_R)
        ELBO[eph] += lbeta(α_R .+ S_R) - lbeta(α_R,R)
    end
    return ELBO, log_λ, (log_θ_R, log_θ_I...)
end

end


using Cubature
using Statistics

function profile(ϕ, params);

    ϕmid = (ϕ[1] + ϕ[end])/2
    Δη = params.Uj * params.Lj * params.TwoΩ / params.g

    U_function(x)  =  params.Uj * sech.((x .- ϕmid) / params.Lj).^2
    LHS(x) = - (params.TwoΩ .* sin.(x) .* params.a .+ U_function(x).*tan.(x)) .* U_function(x) / params.g

    U =  U_function(ϕ)
    E = zeros((params.Nϕ+1))
    for i in 1:params.Nϕ
        result = pquadrature(LHS, ϕ[end], ϕ[i])
        E[i] =  result[1]
    end
    E = E .- mean(E)

    return U,E
end


# Basic zero-mean Gaussian (uni/multi-variate) Metropolis step.
struct StaticRWM{D <: Union{Normal,MvNormal}} <: Metropolis
  proposal::D
end
function _update(rng::Random.AbstractRNG, rwm::StaticRWM, curr::T, logprob::Function) where T
  cand = curr + rand(rng, rwm.proposal)
  log_acceptance_prob = logprob(cand) - logprob(curr)
  accept = log_acceptance_prob > -Random.randexp(rng)
  draw = accept ? cand : curr
  return (draw, accept)
end


using Distributed, StatsBase, Distributions, PyPlot
using TransitionalMCMC

prop(mu) = rand(MvNormal(mu,[3 0 0; 0 3 0; 0 0 3]))

function target(x)
    return log(pdf(MvNormal([-2, 2, 0], [1 0.5 0.5; 0.5 1 0.5; 0.5 0.5 1]), x))
end

samps, acc = metropolis_hastings_simple(target, prop, [-2,2, 0], 2000, 200)

plt.scatter(samps[:, 1],samps[:,2])

f,ax = plt.subplots(3,3)

Ndims = size(samps,2)

for i =1:Ndims
    for j = 1:Ndims
    ax[i,j].scatter(samps[:, i],samps[:,j])
    end
end


@testset "discreteshapebound" begin
    img = generate_surface(SynthSphere(50), 0.5, [0.5,0.0,0.9])
    Z, p, q = DiscreteShapeBound(smoothness=1000, albedo = 1.0, illumination_direction = [0.5,0.0,0.9])(img)
    @test Z[38,38] ≈ 0.00028515016913445543
    @test Z[1,1] ≈ 9.035191974767473e-5
    @test isapprox(Z[76,76], 0.00034007074121476805; atol = 2e-15)
    @test Z[45,45] ≈ 0.00010731186855283559
    @test Z[50,18] ≈ 0.0003244761469170222
    @test Z[11, 68] ≈ 7.131234669714447e-5
    @test Z[42,42] ≈ 0.0002467784663847426
end


function solve_diagTR(λ, Δ, g̃, δ, ϵ; M = [0.0])
    # λ underestimates the required value λstar such that ||d̃(λstar)|| = δ
    # where d̃(λ) =  -(Δ + λ * M) .\ g̃
    T = eltype(δ)
    M = T.(M)

    if M == [0.0]
        # M = ones(Δ) ;
        M = fill(T.(1.0), size(Δ)[1]);
    end
    @assert (minimum(M) > 0.0)
    λin = T.(λ)
    d̃ = -(Δ .+ λ * M) .\ g̃
    d̃d̃ = d̃⋅d̃
    normd̃ = sqrt(d̃d̃)

    tol1 = sqrt(eps(T))
    tolerance = tol1 * T(100.0)

    # Newton iterations
    iter_nwt = 0
    while (normd̃ >= (δ + tolerance * normd̃)) && (iter_nwt < 40)
        dotd̃ = (Δ .+ λ * M) .\ (M .* d̃)
        Δλ = ((normd̃-δ)/δ) * (d̃d̃/(d̃ ⋅ dotd̃))
        λeps = tol1 / T(100.0)
        λ =  max(λ + Δλ, λin + λeps)
        d̃ = -(Δ .+ λ * M) .\ g̃
        d̃d̃ = d̃⋅d̃
        normd̃ = sqrt(d̃d̃)
        @assert (λ >= 0.0)
        iter_nwt += 1
    end

    return d̃, λ
end


using PyDSTool, DataStructures

name = "Calcium channel model"
pars = Dict{String,Any}(
          "vl"=> -60,
          "vca"=> 120,
          "i" => 0,
          "gl" => 2,
          "gca" => 4,
          "c" => 20,
          "v1" => -1.2,
          "v2" => 18)
vars = Dict{String,Any}(
          "v" => "( i + gl * (vl - v) - gca * 0.5 * (1 + tanh( (v-v1)/v2 )) * (v-vca) )/c",
          "w" => "v-w")
ics = Dict{String,Any}(
          "v" => 0,
          "w" => 0)
tdomain = [0;30]

dsargs = build_ode(name,ics,pars,vars,tdomain)

#Solve the ODE
#d = solve_ode(dsargs)
#using Plots
#plot(d[:t],d[:v])

#Bifurcation Plots
ode = ds[:Generator][:Vode_ODEsystem](dsargs)
ode[:set](pars = Dict("i"=>-220))
ode[:set](ics  = Dict("v"=>-170))
PC = ds[:ContClass](ode)

bif = bifurcation_curve(PC,"EP-C",["i"],
                          max_num_points=450,
                          max_stepsize=2,min_stepsize=1e-5,
                          stepsize=2e-2,loc_bif_points="all",
                          save_eigen=true,name="EQ1",
                          print_info=true,calc_stab=true)

@test length(bif.changes) == 2
#=
using Plots
plot(bif,(:i,:v))
=#

bif = bifurcation_curve(PC,"LP-C",["i","gca"],
                        max_num_points=200,initpoint="EQ1:LP2",
                        max_stepsize=2,min_stepsize=1e-5,
                        stepsize=2e-2,loc_bif_points="CP",
                        save_eigen=true,name="SN1",
                        print_info=true,calc_stab=true,
                        solver_sequence=[:forward,:backward])

@test length(bif.changes) == 0

#=
using Plots
plot(bif,(:i,:gca))
=#


# Using Clenshaw's method to evaulate tensor-product Chebyshev polynomials

@generated function clenshaw_recursion(weights::Array{T,N},x::T) where {T,N}

  first_term = :( p = Array{T}(undef,Base.tail(size(weights))) )

  i_vars = Array{Symbol}(undef,N)
  for i = 1:N
    i_vars[i] = Symbol("i$i")
  end

  inner_term = :( z = zeros(T,size(weights,1)+2);
                  for $(i_vars[1]) in size(weights,1):-1:1;
                    z[$(i_vars[1])] = weights[$(i_vars...)]+2*x*z[$(i_vars[1])+1]-z[$(i_vars[1])+2];
                  end;
                  p[$(i_vars[2:N]...)] = z[1]-x*z[2]
                  )

  outer_term = inner_term
  for i = N:-1:2
    outer_term = :( for $(i_vars[i]) = 1:size(weights,$i);
                      $outer_term;
                    end
                    )
  end

  final = :( $first_term;
             $outer_term;
             return p
             )

  return final

end

function clenshaw_evaluate(weights::Array{T,N},x::Array{T,1},domain=[ones(T,1,N);-ones(T,1,N)]) where {T <: AbstractFloat,N}

  for i = 1:N
    normalized_x = normalize_node(x[i],domain[:,i])
    weights = clenshaw_recursion(weights,normalized_x)
  end

  return weights[1]

end

function clenshaw_evaluate(weights::Array{T,N},domain=[ones(T,1,N);-ones(T,1,N)]) where {T <: AbstractFloat,N}

  function goo(x::Array{T,1}) where {T <: AbstractFloat}

    return clenshaw_evaluate(weights,x,domain)

  end

  return goo

end


module TestACSetViews
using Test

using Catlab: @present, generator
using Catlab.Theories: compose, id
using Catlab.Graphs
using Catlab.Graphs.BasicGraphs: TheoryGraph
using Catlab.CSetDataStructures
using Catlab.CategoricalAlgebra.ACSetViews
using Catlab.CategoricalAlgebra.CSets

using Tables

# Basic Indexing and Special Methods
####################################

@present TheoryLocatedGraph <: TheoryGraph begin
  Num::Data
  x::Attr(V,Num)
  y::Attr(V,Num)
end

const LocatedGraph = ACSetType(TheoryLocatedGraph)

g = @acset LocatedGraph{Float64} begin
  V = 5
  E = 4
  src = [1,1,1,1]
  tgt = [2,3,4,5]
  x = [0.,1.,0,-1.,0]
  y = [0.,0.,1.,0.,-1.]
end

edges = ACSetView(g,:E)

upward_edges = @select_where(edges, src.y <= tgt.y)
@test parts(upward_edges) == [1,2,3]

@test upward_edges[1,:src] == 1.
@test upward_edges[1,[:src,:x]] == 0.
@test upward_edges.src == [1,1,1]
@test upward_edges[:,[:tgt,:x]] == [1.,0.,-1.]

@test @compute_prop(upward_edges, (src.x - tgt.x)^2 + (src.y - tgt.y)^2) == [1,1,1]

# Tables.jl Interface
#####################

vertices = ACSetView(g,:V)

@test Tables.istable(vertices)
@test Tables.columnaccess(vertices)
@test Tables.columns(vertices) == vertices
@test Tables.columnnames(vertices) == (:x,:y)
@test Tables.getcolumn(vertices,:x)[2] == 1.
@test Tables.getcolumn(vertices,2,:x) == 1.

# Morphisms
###########

g, h = WeightedGraph{Float64}(2), WeightedGraph{Float64}(4)
add_edge!(g, 1, 2, weight=2.0)
add_edges!(h, [1,2,3], [2,3,4], weight=[1.0,2.0,3.0])
α = ACSetTransformation((V=[2,3], E=[2]), g, h)
αE = ACSetViewMorphism(α, :E)

@test is_natural(αE)
@test αE.f == α.components[:E]

end


# Tolerance
function truncated_svd(A::AbstractMatrix{T}, truncation::Real) where T <: Number
    truncation = min(truncation, one(T))
    U, S, V = svd(A)
    r = vec(S .> truncation*maximum(S))
    U = U[:, r]
    S = S[r]
    V = V[:, r]
    return U, S, V
end

# Explicit rank
function truncated_svd(A::AbstractMatrix{T}, truncation::Int) where T <: Number
    U, S, V = svd(A)
    r = [((i <= truncation && S[i] > zero(T)) ? true : false) for i in 1:length(S)]
    U = U[:, r]
    S = S[r]
    V = V[:, r]
    return U, S, V
end

## Cost functions
# Pareto functions
G(opt::AbstractKoopmanAlgorithm) = g(f) = first(f)

# Evaluate F
function F(opt::AbstractKoopmanAlgorithm)
    f(x, A, y::AbstractArray) = [norm(y .- A*x, 2)] # explicit
    return f
end

# General method with inputs
function (x::AbstractKoopmanAlgorithm)(X::AbstractArray, Y::AbstractArray, U::AbstractArray, B::AbstractArray)
    K, _ = x(X, Y-B*U)
    return (K, B)
end

"""
$(TYPEDEF)

Approximates the Koopman operator `K` based on

```julia
K = Y / X
```

where `Y` and `X` are data matrices. Returns a  `Eigen` factorization of the operator.

# Fields
$(FIELDS)

# Signatures
$(SIGNATURES)
"""
mutable struct DMDPINV <: AbstractKoopmanAlgorithm end;

# Fast but more allocations
function (x::DMDPINV)(X::AbstractArray, Y::AbstractArray)
    K = Y / X
    return (eigen(K),[])
 end

# DMDC
function (x::DMDPINV)(X::AbstractArray, Y::AbstractArray, U::AbstractArray)
    nx, m = size(X)
    nu, m = size(U)

    K̃ = Y / [X;U]
    K = K̃[:, 1:nx]
    B = K̃[:, nx+1:end]

    return (eigen(K), B)
end

"""
$(TYPEDEF)

Approximates the Koopman operator `K` based on the singular value decomposition
of `X` such that:

```julia
K = Y*V*Σ*U'
```

where `Y` and `X = U*Σ*V'` are data matrices. The singular value decomposition is truncated via
the `truncation` parameter, which can either be an `Int` indiciating an index based truncation or a `Real`
indiciating a tolerance based truncation. Returns a `Eigen` factorization of the operator.

# Fields
$(FIELDS)

# Signatures
$(SIGNATURES)
"""
mutable struct DMDSVD{T} <: AbstractKoopmanAlgorithm where T <: Number
    """Indiciates the truncation"""
    truncation::T
end;

DMDSVD() = DMDSVD(0.0)

# Slower but fewer allocations
function (x::DMDSVD{T})(X::AbstractArray, Y::AbstractArray) where T <: Real
    U, S, V = truncated_svd(X, x.truncation)
    xone = one(eltype(X))
    # Computed the reduced operator
    Sinv = Diagonal(xone ./ S)
    B = Y*V*Sinv
    Ã = U'B
    # Compute the modes
    λ, ω = eigen(Ã)
    φ = B*ω
    return (Eigen(λ, φ),[])
end

# DMDc
function (x::DMDSVD{T})(X::AbstractArray, Y::AbstractArray, U::AbstractArray) where T <: Real
    nx, m = size(X)
    nu, m = size(U)
    # Input space svd
    Ũ, S̃, Ṽ = truncated_svd([X;U], x.truncation)
    # Output space svd
    Û, Ŝ, V̂ = svd(Y)

    # Split the svd
    U₁, U₂ = Ũ[1:nx,:], Ũ[nx+1:end,:]

    xone = one(eltype(X))
    # Computed the reduced operator
    C = Y*Ṽ*Diagonal(xone ./ S̃ ) # Common submatrix
    # We do not project onto a reduced subspace here.
    # This would mess up our initial conditions, since sometimes we have
    # x1->x2, x2->x1
    Ã = Û'C*U₁'Û
    B̃ = C*U₂'
    # Compute the modes
    λ, ω = eigen(Ã)
    φ = C*U₁'Û*ω
    return (Eigen(λ, φ), B̃)
end

"""
$(TYPEDEF)

Approximates the Koopman operator `K` based on the singular value decomposition
of `X` and `Y` such that:

```julia
K₁ = Y*V₁*Σ₁*U₁'
K₂ = X*V₂*Σ₂*U₂'
K = (K₁*inv(K₂))^(0.5)
```

where `Y = U₂*Σ₂*V₂'` and `X = U₁*Σ₁*V₁'` are data matrices. The singular value decomposition is truncated via
the `truncation` parameter, which can either be an `Int` indiciating an index based truncation or a `Real`
indiciating a tolerance based truncation. Returns a `Eigen` factorization of the operator.

# Fields
$(FIELDS)

# Signatures
$(SIGNATURES)
"""
mutable struct fbDMD{T} <: AbstractKoopmanAlgorithm where T <: Number
    """Indiciates the truncation"""
    truncation::T
end;
fbDMD() = fbDMD(0.0)
function (x::fbDMD{T})(X::AbstractArray, Y::AbstractArray) where T <: Real
    U₁, S₁, V₁ = truncated_svd(X, x.truncation)
    xone = one(eltype(X))
    # Computed the reduced operator
    Sinv₁ = Diagonal(xone ./ S₁)
    B₁ = Y*V₁*Sinv₁
    A₁ = U₁'B₁  # forward DMD

    U₂, S₂, V₂ = truncated_svd(Y, x.truncation)
    yone = one(eltype(Y))
    # Computed the reduced operator
    Sinv₂ = Diagonal(yone ./ S₂)
    B₂ = X*V₂*Sinv₂
    A₂ = U₂'B₂  # Backward DMD

    Ã = sqrt(A₁/A₂)
    # Compute the modes
    λ, ω = eigen(Ã)
    φ = B₁*ω
    return Eigen(λ, φ)
end
"""
$(TYPEDEF)

Approximates the Koopman operator `K` with the algorithm `alg` over the rank-reduced data
matrices `Xᵣ = X Qᵣ` and `Yᵣ = Y Qᵣ`, where `Qᵣ` originates from the singular value decomposition of
the joint data `Z = [X; Y]`. Based on [this paper](http://cwrowley.princeton.edu/papers/Hemati-2017a.pdf).

If `rtol` ∈ (0, 1) is given, the singular value decomposition is reduced to include only
entries bigger than `rtol*maximum(Σ)`. If `rtol` is an integer, the reduced SVD up to `rtol` is used
for computation.

# Fields
$(FIELDS)

# Signatures
$(SIGNATURES)
"""
mutable struct TOTALDMD{R, A} <: AbstractKoopmanAlgorithm where {R <: Number, A <: AbstractKoopmanAlgorithm}
    truncation::R
    alg::A
end

TOTALDMD() = TOTALDMD(0.0, DMDPINV())

function (x::TOTALDMD)(X::AbstractArray, Y::AbstractArray)
    _ , _, Q = truncated_svd([X; Y], x.truncation)
    return x.alg(X*Q, Y*Q)
end

function (x::TOTALDMD)(X::AbstractArray, Y::AbstractArray, U::AbstractArray)
    _ , _, Q = truncated_svd([X; Y], x.truncation)
    return x.alg(X*Q, Y*Q, U*Q)
end

function (x::TOTALDMD)(X::AbstractArray, Y::AbstractArray, U::AbstractArray, B::AbstractArray)
    _ , _, Q = truncated_svd([X; Y], x.truncation)
    K, _ = x.alg(X*Q, (Y-B*U)*Q)
    return (K, B)
end


"""
```
  lsq_constrsparsereg(
      X         :: AbstractMatrix{T},
      y         :: AbstractVector{T},
      ρ         :: Union{AbstractVector, Number} = zero(T);
      Aeq       :: AbstractMatrix = zeros(T, 0, size(X, 2)),
      beq       :: Union{AbstractVector, Number} = zeros(T, size(Aeq, 1)),
      Aineq     :: AbstractMatrix = zeros(T, 0, size(X, 2)),
      bineq     :: Union{AbstractVector, number} = zeros(T, size(Aineq, 1)),
      obswt     :: AbstractVector = ones(T, length(y)),
      penwt     :: AbstractVector = ones(T, size(X, 2)),
      warmstart :: Bool = false,
      solver = ECOSSolver(maxit=10e8, verbose=0)
    )
```
Fit constrained lasso at fixed tuning parameter value(s) by minimizing
    `0.5sumabs2(√obswt .* (y - X * β)) + ρ * sumabs(penwt .* β)`
subject to linear constraints, using `Convex.jl`.

### Arguments
- `X`       : predictor matrix.
- `y`       : response vector.
- `ρ`       : tuning parameter. Can be a number or a list of numbers. Default 0.

### Optional arguments
- `Aeq`       : equality constraint matrix.
- `beq`       : equality constraint vector.
- `Aineq`     : inequality constraint matrix.
- `bineq`     : inequality constraint vector.
- `obswt`     : observation weights. Default is `[1 1 1 ... 1]`.
- `penwt`     : predictor penalty weights. Default is `[1 1 1 ... 1]`.
- `warmstart` :
- `solver`    : a solver Convex.jl supports. Default is ECOS.
              Note that Mosek and Gurobi are more robust than ECOS. Unlike ECOS or
              SCS, both Mosek and Gurobi require a license (free for academic
              use). <http://convexjl.readthedocs.io/en/latest/solvers.html>

### Returns
- `β`       : estimated coefficents.
- `objval`  : optimal objective value.
- `problem` : Convex.jl problem.
"""

function lsq_constrsparsereg(
    X::AbstractMatrix{T},
    y::AbstractVector{T},
    ρ::Union{AbstractVector, Number} = zero(T);
    Aeq::AbstractMatrix = zeros(T, 0, size(X, 2)),
    beq::Union{AbstractVector, Number} = zeros(T, size(Aeq, 1)),
    Aineq::AbstractMatrix = zeros(T, 0, size(X, 2)),
    bineq::Union{AbstractVector, Number} = zeros(T, size(Aineq, 1)),
    obswt::AbstractVector = ones(T, length(y)),
    penwt::AbstractVector = ones(T, size(X, 2)),
    warmstart::Bool = false,
    solver = ECOSSolver(maxit=10e8, verbose=0)
    ) where T

    n, p = size(X)

    β̂ = zeros(T, p, length(ρ))
    optval_vec = zeros(length(ρ))
    prob_vec = Convex.Problem[]

    β = Variable(p)
    loss = (1//2)sumsquares(sqrt(obswt) .* (y - X * β)) # loss term
    pen  = dot(penwt, abs(β)) # penalty term
    eqconstr   = Aeq * β == beq
    ineqconstr = Aineq * β <= bineq

    for i in 1:length(ρ)

      ρi = ρ[i]
      if ρi == typemax(typeof(ρi))
          problem = minimize(pen, eqconstr, ineqconstr)
      elseif ρi ≤ 0
          problem = minimize(loss, eqconstr, ineqconstr)
      else
          problem = minimize(loss + ρi * pen, eqconstr, ineqconstr)
      end

      if warmstart
          solve!(problem, solver; warmstart = i > 1? true : false)
          push!(prob_vec, problem)
          optval_vec[i] = problem.optval
      else
          solve!(problem, solver)
          push!(prob_vec, problem)
          optval_vec[i] = problem.optval
      end

      if length(ρ) == 1
        return β.value, problem.optval, problem
      end

      β̂[:, i] = β.value

    end # end of for loop

    return β̂, optval_vec, prob_vec

end # end of function


using GLAbstraction, Colors, GeometryTypes, GLVisualize, Reactive

if !isdefined(:runtests) # We use the examples to generate docs and tests
    window = glscreen()
    timesignal = bounce(linspace(0.0, 1.0, 360))
end

n = 400 # The number of points per line
nloops = 20 # The number of loops
# The scalar parameter for each line
TL = linspace(-2f0 * pi, 2f0 * pi, n)
# We create a list of positions and connections, each describing a line.
# We will collapse them in one array before plotting.
xyz    = Point3f0[]
colors = RGBA{Float32}[]

# creates some distinguishable colors from which we can sample for each line
base_colors1 = distinguishable_colors(nloops, RGB{Float64}(1,0,0))
# Create each line one after the other in a loop
for i=1:nloops
    append!(xyz, [Point3f0(sin(t), cos((2 + .02 * i) * t), cos((3 + .03 * i) * t)) for t in TL])
    unique_colors = base_colors1[i]
    hsv = HSV(unique_colors)
    color_palette = map(x->RGBA{Float32}(x, 1.0), sequential_palette(hsv.h, n, s=hsv.s))
    append!(colors, color_palette)
end

# map comes from Reactive.jl and allows you to map any Signal to another.
# In this case we create a rotation matrix from the timesignal signal.

rotation = map(timesignal) do t
    rotationmatrix_z(Float32(t*2pi)) # -> 4x4 Float32 rotation matrix
end

lines3d = visualize(xyz, :lines, color=colors, model=rotation)

_view(lines3d, window, camera=:perspective)

if !isdefined(:runtests)
    renderloop(window)
end


using Test, LinearAlgebra
using QBase

@testset "./src/constructors/states.jl" begin

@testset "pure_state()" begin
    @testset "valid inputs" begin
        ρ = pure_state(Ket([im,0]))
        @test ρ isa State{Complex{Int64}}
        @test rank(ρ) == 1
        @test ρ == [1 0;0 0]

        ρ = pure_state(Bra([0.,-1im]))
        @test ρ isa State{Complex{Float64}}
        @test rank(ρ) == 1
        @test ρ == [0 0;0 1]

        ρ = pure_state([1,1]/sqrt(2))
        @test ρ isa State{Float64}
        @test rank(ρ) == 1
        @test ρ ≈ [1 1;1 1]/2

        ρ = pure_state(([1,-im]/sqrt(2))')
        @test ρ isa State{Complex{Float64}}
        @test rank(ρ) == 1
        @test ρ ≈ [1 im;-im 1]/2
    end

    @test_throws DomainError pure_state([1;1])
end

@testset "mixed_state()" begin
    @testset "strongly typed inputs" begin
        ρ_mix = mixed_state(
            Probabilities([0.7,0.2,0.1]),
            State.([[1 0 0;0 0 0;0 0 0],[0 0 0;0 0 0;0 0 1],[0 0 0;0 1 0;0 0 0]])
        )
        @test ρ_mix == [0.7 0 0;0 0.1 0;0 0 0.2]
        @test ρ_mix isa State{Float64}
        @test rank(ρ_mix) == 3
    end

    @testset "loosely typed inputs" begin
        ρ_mix = mixed_state([0.7,0.2,0.1],[[1 0 0;0 0 0;0 0 0],[0 0 0;0 0 0;0 0 1],[0 0 0;0 1 0;0 0 0]])
        @test ρ_mix == [0.7 0 0;0 0.1 0;0 0 0.2]
        @test ρ_mix isa State{Float64}
        @test rank(ρ_mix) == 3

        ρ_mix = mixed_state(
            Probabilities([1/2, 1/2]),
            [0.5*[1 im;-im 1], [1 0;0 0]]
        )
        @test ρ_mix == [0.75 0.25im;-0.25im 0.25]
        @test ρ_mix isa State{ComplexF64}
        @test rank(ρ_mix) == 2

        @test_throws DomainError mixed_state([0.7,0.2], [[1 0;0 0],[0 0;0 1]])
        @test_throws DomainError mixed_state([0.7,0.3], [[1 0;0 .1],[0 0;0 1]])
    end
end

@testset "trine_qubit_states()" begin
    trine_states = trine_qubit_states()
    @test isapprox(trine_states, mirror_symmetric_qubit_states(π/3), atol=1e-7)
    @test trine_states isa Vector{State{Float64}}
end

@testset "mirror_symmetric_qubit_states()" begin
    @testset "range over θ" begin
        for θ in 0:0.1π:π/2
            @test mirror_symmetric_qubit_states(θ) isa Vector{State{Float64}}
        end
    end

    @test_throws DomainError mirror_symmetric_qubit_states(π)

    @test isapprox(mirror_symmetric_qubit_states(π/2),[[1 0; 0 0], [0 0; 0 1], [0 0; 0 1]], atol=1e-7)
end

@testset "bloch_qubit_state()" begin
    @test [1 0;0 0] == bloch_qubit_state(0,0,1)
    @test 0.5*[1 -im;im 1] == bloch_qubit_state(0,1,0)
    @test 0.5*[1 1;1 1] == bloch_qubit_state(1,0,0)
    @test 0.5*[1 0;0 1] == bloch_qubit_state(0,0,0)

    @test bloch_qubit_state(0,0,1) == bloch_qubit_state(0,0)
    @test isapprox(bloch_qubit_state(0,1,0),bloch_qubit_state(π/2,π/2),atol=1e-7)
    @test isapprox(bloch_qubit_state(1,0,0),bloch_qubit_state(π/2,0),atol=1e-7)

    @test bloch_qubit_state(0,0,1) isa State{Complex{Float64}}
    @test_throws DomainError bloch_qubit_state(1,1,1)
end

@testset "sic_qubit_states()" begin
    sic_states = sic_qubit_states()

    @test sic_states isa Vector{State{Complex{Float64}}}

    θ_sic = 2*acos(1/sqrt(3))

    @test sic_states[1] ≈ bloch_qubit_state(0, 0)
    @test sic_states[2] ≈ bloch_qubit_state(θ_sic, 0)
    @test sic_states[3] ≈ bloch_qubit_state(θ_sic, 2π/3)
    @test sic_states[4] ≈ bloch_qubit_state(θ_sic, 4π/3)
end

@testset "bb84_qubit_states()" begin
    bb84_states = bb84_qubit_states()
    @test bb84_states isa Vector{State{Float64}}

    @test isapprox(bb84_states, planar_symmetric_qubit_states(4), atol=1e-7)
end

@testset "computational_basis_states()" begin
    @test computational_basis_states(3) == [[1 0 0;0 0 0;0 0 0],[0 0 0;0 1 0;0 0 0],[0 0 0;0 0 0;0 0 1]]
    @test computational_basis_states(2) == [[1 0;0 0],[0 0;0 1]]
end

@testset "bell_states()" begin
    states = bell_states()

    @test states isa Vector{State{Float64}}
    @test states[1] ≈ 1/2*[1 0 0 1;0 0 0 0;0 0 0 0;1 0 0 1]
    @test states[2] ≈ 1/2*[1 0 0 -1;0 0 0 0;0 0 0 0;-1 0 0 1]
    @test states[3] ≈ 1/2*[0 0 0 0;0 1 1 0;0 1 1 0;0 0 0 0]
    @test states[4] ≈ 1/2*[0 0 0 0;0 1 -1 0;0 -1 1 0;0 0 0 0]
end

@testset "generalized_bell_states()" begin
    @testset "qubit bell states" begin
        states = generalized_bell_states(2)

        @test states isa Vector{State{Complex{Float64}}}
        @test isapprox(states, bell_states(),atol=1e-7)
    end
end

@testset "planar_symmetric_qubit_states()" begin
    @test isapprox(planar_symmetric_qubit_states(2), [[1 0;0 0],[0 0;0 1]], atol=1e-7)
    @test isapprox(planar_symmetric_qubit_states(3), trine_qubit_states(),atol=1e-7)

    @testset "scanning over cases verifying right number are computed" begin
        for n in 5:50
            qubits = planar_symmetric_qubit_states(n)

            @test length(qubits) == n
            @test qubits isa Vector{State{Float64}}
        end
    end
end

end


export demo

using Gadfly

function demo()
# Setup
    x = collect(1.0:1.0:20)
    u = x.*0.0
    s = x
    v = x.*0.0
    t = x
    w = x.*0.0

    for i in 6:10
        u[i] = 1.0
    end
    for i in 11:14
        u[i] = -0.2 * i + 3.0
    end
    for i in 11:14
        v[i] = -0.2
    end

    p = collect(1.0:0.2:20)
    r = p.*0.0
    tol = 1e-7
    i = 0
    for pi in p
        i += 1
        if pi >= (6.0 - tol) && pi <= (10.0 + tol)
            r[i] = 1.0
        end
        if pi > (10.0 + tol) && pi < (15.0 - tol)
            r[i] = -0.2 * p[i] + 3.0
        end
    end
####
    interpolate(x, u, RK_W3())

    cond = get_cond()
    ε = get_epsilon()
    println("cond = $cond, ε = $ε")

    σ = evaluate(p)

    δ = σ - r
    rms = round(norm(δ)/sqrt(length(δ)); digits=2)
    println("rms = $rms")

    set_default_plot_size(13cm, 13cm)
    plt = plot(layer(x = x, y = u, Geom.point, Theme(default_color=colorant"orange")),
          layer(x = p, y = r, Geom.line, Theme(default_color=colorant"red")),
          layer(x = p, y = σ, Geom.line, Theme(default_color=colorant"blue")),
          Scale.y_continuous(minvalue=-0.5, maxvalue=1.5),
          Guide.manual_color_key("Legend", ["Points", "True", "Spline"], ["orange", "red", "blue"]),
          Guide.xlabel("knots, points"), Guide.ylabel("f, σ"),
          Guide.title("Fig.1a"))
    draw(SVG("plot_W3.svg", 13cm, 13cm), plt)

    σ = evaluate(p, 1)
    plt = plot(layer(x = x, y = u, Geom.point, Theme(default_color=colorant"orange")),
          layer(x = p, y = σ, Geom.line, Theme(default_color=colorant"blue")),
          Scale.y_continuous(minvalue=-0.5, maxvalue=1.5),
          Guide.manual_color_key("Legend", ["Points", "Spline 1st derivative"], ["orange", "blue"]),
          Guide.xlabel("knots, points"), Guide.ylabel(raw"σ'"),
          Guide.title("Fig.2a"))
    draw(SVG("plot_W3_d1.svg", 13cm, 13cm), plt)

    σ = evaluate(p, 2)
    plt = plot(layer(x = x, y = u, Geom.point, Theme(default_color=colorant"orange")),
          layer(x = p, y = σ, Geom.line, Theme(default_color=colorant"blue")),
          Scale.y_continuous(minvalue=-0.5, maxvalue=1.5),
          Guide.manual_color_key("Legend", ["Points", "Spline 2nd derivative"], ["orange", "blue"]),
          Guide.xlabel("knots, points"), Guide.ylabel(raw"σ''"),
          Guide.title("Fig.3a"))
    draw(SVG("plot_W3_d2.svg", 13cm, 13cm), plt)

###
      interpolate(x, u, s, v, t, w, RK_W3())

      cond = get_cond()
      ε = get_epsilon()
      println("cond = $cond, ε = $ε")

      σ = evaluate(p)

      δ = σ - r
      rms = round(norm(δ)/sqrt(length(δ)); digits=2)
      println("rms = $rms")

      set_default_plot_size(13cm, 13cm)
      plt = plot(layer(x = x, y = u, Geom.point, Theme(default_color=colorant"orange")),
            layer(x = p, y = r, Geom.line, Theme(default_color=colorant"red")),
            layer(x = p, y = σ, Geom.line, Theme(default_color=colorant"blue")),
            Scale.y_continuous(minvalue=-0.5, maxvalue=1.5),
            Guide.manual_color_key("Legend", ["Points", "True", "Spline"], ["orange", "red", "blue"]),
            Guide.xlabel("knots, points"), Guide.ylabel("f, σ"),
            Guide.title("Fig.1b"))
      draw(SVG("plot_d12_W3.svg", 13cm, 13cm), plt)

      σ = evaluate(p, 1)
      plt = plot(layer(x = x, y = u, Geom.point, Theme(default_color=colorant"orange")),
            layer(x = p, y = σ, Geom.line, Theme(default_color=colorant"blue")),
            Scale.y_continuous(minvalue=-0.5, maxvalue=1.5),
            Guide.manual_color_key("Legend", ["Points", "Spline 1st derivative"], ["orange", "blue"]),
            Guide.xlabel("knots, points"), Guide.ylabel(raw"σ'"),
            Guide.title("Fig.2b"))
      draw(SVG("plot_d12_W3_d1.svg", 13cm, 13cm), plt)

      σ = evaluate(p, 2)
      plt = plot(layer(x = x, y = u, Geom.point, Theme(default_color=colorant"orange")),
            layer(x = p, y = σ, Geom.line, Theme(default_color=colorant"blue")),
            Scale.y_continuous(minvalue=-0.5, maxvalue=1.5),
            Guide.manual_color_key("Legend", ["Points", "Spline 2nd derivative"], ["orange", "blue"]),
            Guide.xlabel("knots, points"), Guide.ylabel(raw"σ''"),
            Guide.title("Fig.3b"))
      draw(SVG("plot_d12_W3_d2.svg", 13cm, 13cm), plt)

end


using StatsBase
using Test

@testset "ECDF" begin
    x = randn(10000000)
    fnecdf = ecdf(x)
    y = [-1.96, -1.644854, -1.281552, -0.6744898, 0, 0.6744898, 1.281552, 1.644854, 1.96]
    @test isapprox(fnecdf(y), [0.025, 0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 0.95, 0.975], atol=1e-3)
    @test isapprox(fnecdf(1.96), 0.975, atol=1e-3)
    @test fnecdf(y) ≈ map(fnecdf, y)
    @test extrema(fnecdf) == (minimum(fnecdf), maximum(fnecdf)) == extrema(x)
    fnecdf = ecdf([0.5])
    @test fnecdf([zeros(5000); ones(5000)]) == [zeros(5000); ones(5000)]
    @test extrema(fnecdf) == (minimum(fnecdf), maximum(fnecdf)) == (0.5, 0.5)
end


#using DelimitedFiles # To be able to load .txt files

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

"""
    conversion_factor_arcsec_to_sma

Multiplicative conversion factor from arcseconds to mpc.
"""
const conversion_factor_arcsec_to_sma = d_BH * 2*pi/(360*60*60)

"""
    arcsec_to_mpc(a::Float64)

Converts a semi-major axis given in arcseconds to an semi-major axis given in mpc.
"""
function arcsec_to_mpc(a::Float64)
    return a * conversion_factor_arcsec_to_sma
end

"""
    get_dataCluster()

Recovers the S-cluster data from INPUTCLUSTER.
"""
function get_dataCluster()
    dataCluster, headerCluster = readdlm(INPUTCLUSTER,header=true) # Reading the data file with the header
    return dataCluster # Output the entire data
end

"""
    dataCluster

Table of S-cluster data
"""
const dataCluster = get_dataCluster() # Reading the input data from the Cluster

"""
    nbCluster

Number of S-cluster stars used in the code
"""
const nbCluster = size(dataCluster)[1] # Number of stars in cluster

#################################################
# Getting the relevant physical quantities of the S-cluster
##################################################

"""
    tabaCluster

Table of S-cluster SMAs
"""
const tabaCluster = SVector{nbCluster,Float64}([arcsec_to_mpc(dataCluster[iCluster,2]) for iCluster=1:nbCluster]) # The sma a is stored in column 2 (mpc)

"""
    tabjCluster

Table of S-cluster reduced angular momenta.
"""
const tabjCluster = SVector{nbCluster,Float64}([ecc_to_j(dataCluster[iCluster,3]) for iCluster=1:nbCluster]) # The eccentricity e is stored in column 3

"""
    tabTCluster

Table of S-cluster main-sequence ages.
"""
const tabTCluster = SVector{nbCluster,Float64}([dataCluster[iCluster,4] for iCluster=1:nbCluster]) # The main-sequence age is stored in column 4 (in Myr)

"""
    tabdjCluster

Table of S-cluster SMA uncertainties.
"""
const tabdjCluster = SVector{nbCluster,Float64}([uncertainty_ecc_to_j(dataCluster[iCluster,3],dataCluster[iCluster,5]) for iCluster=1:nbCluster]) # The eccentricity uncertainty de is stored in column 5

"""
    tabdTpCluster

Table of S-cluster upper main-sequence age uncertainties.
"""
const tabdTpCluster = SVector{nbCluster,Float64}([dataCluster[iCluster,6] for iCluster=1:nbCluster]) # The age up-uncertainty dT_p is stored in column 6

"""
    tabdTmCluster

Table of S-cluster lower main-sequence age uncertainties.
"""
const tabdTmCluster = SVector{nbCluster,Float64}([dataCluster[iCluster,7] for iCluster=1:nbCluster]) # The age down-uncertainty dT_m is stored in column 7


function rand_discr(p; m = 1)
    """
        rand_discr - discrete random generator

          y = rand_discr(p, n);

          y is a random vector of length n drawn from
          a variable X such that
              p(i) = Prob( X=i )

          Copyright (c) 2004 Gabriel Peyré
    """

    # makes sure it sums to 1
    p = p/sum(p)

    n = length(p)
    coin = rand(m)
    cumprob = [0.]
    append!(cumprob, cumsum(p))
    sample = zeros(m)

    for j in 1:n
        ind = findall((coin .> cumprob[j]) & (coin .<= cumprob[j+1]))
        sample[ind] = j
    end
    return sample
end


function jt07a()
#
#  test polynomial suggested by Jenkins and Traub
#
   a = 10.0^(-10);
   z = [.001; .01; .1; .1+a*im; .1-a*im; 1; -10]
   p = poly(z)
   z = [z ones(7)]
   y = [-10  1; 1  1; 0.01  1; 0.001  1; 0.1  3];
   p, PolyZeros(z), PolyZeros(y)
end


"""
Implements toy data from paper
"""
function generate_toy_data(N_train=50)

    X = fill(collect(0.005:0.005:1.), 3)
    Y = fill(Float64[], 3)

    Sl = [5, 1, 2]
    Pl = [200, 0.1, 100]

    conv_ker = (Δ, S, P) ->  S * exp(-P * Δ^2)
    u = t -> sum(k -> (1 / k^2) * cos(2 * k * π * t), 1:4)

    for i in 1:3
        y = fill(0., 200)
        for j in 1:200
            y[j] =  quadgk(τ -> conv_ker(X[i][j] - τ, Sl[i], Pl[i]) * u(τ), 0., X[i][j])[1]
        end
        y = sum(x -> y.^x, 1:3)
        Y[i] = y +  randn(200) * sqrt(0.005 * (sum((y .- (sum(y) / 200)).^2) / 200))
    end

    train_X = fill(Float64[], 3)
    test_X = fill(Float64[], 3)
    train_Y = fill(Float64[], 3)
    test_Y = fill(Float64[], 3)
    for i in 1:3
        mix = randperm(200)
        mixtr = sort(mix[1:N_train])
        mixte = sort(mix[N_train:end])

        train_X[i] = X[i][mixtr]
        test_X[i] = X[i][mixte]
        train_Y[i] = Y[i][mixtr]
        test_Y[i] = Y[i][mixte]
    end

    (Data(train_X, train_Y), Data(test_X, test_Y))
end

function load_weather(N_train)

end


function readinput(filename)
    return map(x -> split(x, " = "), readlines(filename))
end

function masks(mask)
    and, or, x = UInt(0), UInt(0), UInt(0)
    for c in mask
        and, or, x = [and, or, x] .<<= 1
        if c == '0'
            and += 1
        elseif c == '1'
            or += 1
        else
            x += 1
        end
    end
    return (~and, or, x)
end

function allcombinations(addr, floating)
    if floating == 0
        return addr
    else
        attempt = UInt(1)
        while true
            if attempt == floating & attempt
                return cat(allcombinations(addr, floating &= ~attempt), allcombinations(xor(addr, attempt), floating &= ~attempt); dims=1)
            end
            attempt <<= 1
        end
    end
end

function memoryallocation(instructions)
    P1, P2 = Dict(), Dict()
    and = or = x = 0
    for instruction in instructions
        cmd, val = instruction[1], instruction[2]
        if cmd == "mask"
            and, or, x = masks(val)
        else
            addr, val = parse(UInt, match(r"^mem\[(\d+)\]$", cmd)[1]), parse(UInt, val)
            P1[addr] = (val | or) & and
            for loc in allcombinations(addr | or, x)
                P2[loc] = val
            end
        end
    end
    return P1, P2
end

instructions = readinput("d14.in")
P1, P2 = memoryallocation(instructions)

println("P1: ", sum(values(P1)))
println("P2: ", sum(values(P2)))


#=%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Example : Find the distribution of a tracer given:              %
%       (a) the pathways described by A,                          %
%       (b) interior sources and sinks given by dC,               %
%           that best fits observations, Cobs,                    %
%   and (c) inequality constraints on the tracer concentration.   %
%                                                                  %
% Mathematically, minimize J = (C-Cobs)^T W (C-Cobs) subject to    %
%                         AC = d + Gamma u                         %
%  where u is the estimated change in surface concentration.    %
%
% See Supplementary Section 2, Gebbie & Huybers 2011.
% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% =#

using Revise, TMI, GoogleDrive
using PyPlot, PyCall, Test
#, Distributions, LinearAlgebra,  Zygote, ForwardDiff, Optim

TMIversion = "modern_90x45x33_GH10_GH12"
A, Alu, γ, TMIfile, L, B = config_from_nc(TMIversion)

# first guess of change to surface boundary conditions
# ocean values are 0
#u = zerosurfaceboundary(γ)
u = (;surface = zerosurfaceboundary(γ))
uvec = vec(u)

# take synthetic, noisy observations
y, W⁻, ctrue = synthetic_observations(TMIversion,"θ",γ)

# a first guess: observed surface boundary conditions are perfect.
# set surface boundary condition to the observations.
# below surface = 0 % no internal sinks or sources.
#b = getsurfaceboundary(y)
b = (;surface = getsurfaceboundary(y))

# get sample J value
F,G = costfunction_gridded_obs(uvec,Alu,b,u,y,W⁻,γ)
fg!(F,G,x) = costfunction_gridded_obs!(F,G,x,Alu,b,u,y,W⁻,γ)

fg(x) = costfunction_gridded_obs(x,Alu,b,u,y,W⁻,γ)
f(x) = fg(x)[1]
J₀,gJ₀ = fg(uvec)

#### gradient check ###################
# check with forward differences
ϵ = 1e-3
ii = rand(1:sum(γ.wet[:,:,1]))
println("Location for test =",ii)
δu = copy(uvec); δu[ii] += ϵ
∇f_finite = (f(δu) - f(uvec))/ϵ
println(∇f_finite)

fg!(J₀,gJ₀,(uvec+δu)./2) # J̃₀ is not overwritten
∇f = gJ₀[ii]
println(∇f)

# error less than 10 percent?
println("Percent error ",100*abs(∇f - ∇f_finite)/abs(∇f + ∇f_finite))
#### end gradient check #################

# filter the data with an Optim.jl method
iterations = 15
out = steadyclimatology(uvec,fg!,iterations)

# reconstruct by hand to double-check.
ũ = unvec(u,out.minimizer)

# apply to the boundary conditions
b̃ = adjustboundarycondition(b,ũ)

# reconstruct tracer map
c₀ = steadyinversion(Alu,b,γ)
c̃  = steadyinversion(Alu,b̃,γ)

Δc̃ = c̃ - ctrue
Δc₀ = c₀ - ctrue

# plot the difference
level = 15 # your choice 1-33
depth = γ.depth[level]
cntrs = -1:0.05:1
label = "Optimized misfit: Δc̃"
# Help: needs work with continents and labels
planviewplot(Δc̃, depth, cntrs, titlelabel=label)

label = "First guess misfit: Δc₀"
planviewplot(Δc₀, depth, cntrs, titlelabel=label)

# future step: box minimization to eliminate unreasonable temperatures.
# should be made more generic for other tracers
# From MATLAB:
#lbT = -2.*ones(Nsfc,1); % temperature lower bound: can not freeze.
#ubT = 40.*ones(Nsfc,1); % ad-hoc temperature upper bound: 40 C.


"""
Compute p(x) where p is the unique polynomial of degree length(xi),
such that p(x[i]) = y[i]) for all i.

    ti: interpolation nodes
    xi: interpolation values
    t:  evaluation point
    x:  evaluation value
"""
function aitken_neville(ti::Vector{TT}, xi::Matrix, t::TT, x::Vector) where {TT}
    @assert length(ti) == size(xi,2)
    @assert length(x)  == size(xi,1)

    for j in axes(t)
        for i in 1:(length(t)-j)
            for k in axes(x,1)
                xi[k,i] = xi[k,i+1] + (xi[k,i] - xi[k,i+1]) * (ti[i+j] - t) / (ti[i+j] - ti[i])
            end
        end
    end
    for k in eachindex(x)
        x[k] = xi[k,1]
    end
end

"""
Euler extrapolation method with arbitrary order p.

    v:  function to compute vector field
    t₀: initial time
    t₁: final   time
    x₀: initial value
    x₁: final   value
    s:  number of interpolations (order p=s+1)

# TODO This is probably broken!
"""
function euler_extrapolation(v::Function, t₀::TT, t₁::TT, x₀::Vector, x₁::Vector, s::Int) where {TT}
    @assert size(x₀) == size(x₁)

    F   = collect(1:(s+1))
    Δt  = t₁ - t₀
    σ   = Δt ./ F
    pts = repmat(x₀, 1, s+1)

    local xᵢ = zero(x₀)
    local vᵢ = zero(x₀)

    for i in 1:(s+1)
        for j in 1:(F[i]-1)
            tᵢ  = t₀ + σ[i]
            for k in axes(pts,1)
                xᵢ[k] = pts[k,i]
            end
            v(tᵢ, xᵢ, vᵢ)
            for k in axes(pts,1)
                pts[k,i] += σ[i] * vᵢ[k]
            end
        end
    end

    aitken_neville(σ, pts, zero(TT), x₁)
end

"""
Midpoint extrapolation method with arbitrary order p.

    v:  function to compute vector field
    t₀: initial time
    t₁: final   time
    x₀: initial value
    x₁: final   value
    s:  number of interpolations (order p=2s+2)
"""
function midpoint_extrapolation(v::Function, t₀::TT, t₁::TT, x₀::Vector, x₁::Vector, s::Int) where {TT}
    @assert size(x₀) == size(x₁)

    local F   = 2*collect(1:(s+1))
    local Δt  = t₁ - t₀
    local σ   = Δt ./ F
    local σ²  = σ.^2
    local pts = zeros(eltype(x₀), length(x₀), s+1)

    local xᵢ₁= zero(x₀)
    local xᵢ₂= zero(x₀)
    local xᵢₜ= zero(x₀)
    local vᵢ = zero(x₀)
    local v₀ = zero(x₀)

    v(t₀, x₀, v₀)

    for i in 1:s+1
        tᵢ   = t₀ + σ[i]
        xᵢ₁ .= x₀
        xᵢ₂ .= x₀ + σ[i] .* v₀
        for j in 1:(F[i]-1)
            v(tᵢ, xᵢ₂, vᵢ)
            xᵢₜ .= xᵢ₁ + 2σ[i] .* vᵢ
            xᵢ₁ .= xᵢ₂
            xᵢ₂ .= xᵢₜ
        end
        for k in axes(pts,1)
            pts[k,i] += xᵢ₂[k]
        end
    end

    aitken_neville(σ², pts, zero(TT), x₁)
end

"""
Midpoint extrapolation method with arbitrary order p.

    v:  function to compute vector field
    f:  function to compute force  field
    t₀: initial time
    t₁: final   time
    q₀: initial positions
    p₀: initial momenta
    q₁: final   positions
    p₁: final   momenta
    s:  number of interpolations (order p=2s+2)
"""
function midpoint_extrapolation(v::Function, f::Function, t₀::TT, t₁::TT, q₀::Vector, q₁::Vector, p₀::Vector, p₁::Vector, s::Int) where {TT}
    @assert size(q₀) == size(q₁) == size(p₀) == size(p₁)

    local F   = 2*collect(1:(s+1))
    local Δt  = t₁ - t₀
    local σ   = Δt ./ F
    local σ2  = σ.^2

    local qts = zeros(eltype(q₀), length(q₀), s+1)
    local pts = zeros(eltype(p₀), length(p₀), s+1)

    local qᵢ₁= zero(q₀)
    local qᵢ₂= zero(q₀)
    local qᵢₜ= zero(q₀)

    local pᵢ₁= zero(p₀)
    local pᵢ₂= zero(p₀)
    local pᵢₜ= zero(p₀)

    local v₀ = zero(q₀)
    local vᵢ = zero(q₀)

    local f₀ = zero(p₀)
    local fᵢ = zero(p₀)

    v(t₀, q₀, p₀, v₀)
    f(t₀, q₀, v₀, f₀)

    for i in 1:(s+1)
        tᵢ   = t₀ + σ[i]
        qᵢ₁ .= q₀
        qᵢ₂ .= q₀ + σ[i] .* v₀
        pᵢ₁ .= p₀
        pᵢ₂ .= p₀ + σ[i] .* f₀
        for j in 1:(F[i]-1)
            v(tᵢ, qᵢ₂, pᵢ₂, vᵢ)
            f(tᵢ, qᵢ₂, vᵢ,  fᵢ)
            qᵢₜ .= qᵢ₁ + 2σ[i] .* vᵢ
            qᵢ₁ .= qᵢ₂
            qᵢ₂ .= qᵢₜ
            pᵢₜ .= pᵢ₁ + 2σ[i] .* fᵢ
            pᵢ₁ .= pᵢ₂
            pᵢ₂ .= pᵢₜ
        end
        for k in axes(qts,1)
            qts[k,i] += qᵢ₂[k]
        end
        for k in axes(pts,1)
            pts[k,i] += pᵢ₂[k]
        end
    end

    aitken_neville(σ2, qts, zero(TT), q₁)
    aitken_neville(σ2, pts, zero(TT), p₁)
end


function compute_log_likelihood(n::Int, p::Int, v::Matrix{T}, F::Array{T, 3}, valid_insts::Vector{Int}) where T
    if p == 1
        return compute_log_likelihood_univariate(n, v, F, valid_insts)
    else
        return compute_log_likelihood_multivariate(n, p, v, F, valid_insts)
    end
end

function compute_log_likelihood(n::Int, p::Int, v::Vector{T}, F::Vector{T}, valid_insts::Vector{Int}) where T
    log_likelihood::T = n*log(2*pi)/2
    @inbounds for t in valid_insts
        log_likelihood += 0.5 * (log(F[t]) + (v[t]^2)/F[t])
    end
    return log_likelihood
end

function compute_log_likelihood_univariate(n::Int, v::Matrix{T}, F::Array{T, 3}, valid_insts::Vector{Int}) where T
    log_likelihood::T = n*log(2*pi)/2
    @inbounds for t in valid_insts
        log_likelihood += 0.5 * (log(F[1, 1, t]) + (v[t, 1]^2)/F[1, 1, t])
    end
    return log_likelihood
end

function compute_log_likelihood_multivariate(n::Int, p::Int, v::Matrix{T}, F::Array{T, 3}, valid_insts::Vector{Int}) where T
    log_likelihood::T = n*p*log(2*pi)/2
    @inbounds @views for t in valid_insts
        log_likelihood += 0.5 * (logdetF(F, t) + (v[t, :]' * invertF(F, t) * v[t, :]))
    end
    return log_likelihood
end

function get_log_likelihood_params(model::StateSpaceModel{T},
                                   filter_type::DataType) where T
    error(filter_type , " not implemented") # Returns an error if it cannot
                                            # find a specialized get_log_likelihood_params
end

function valid_instants(model::StateSpaceModel{T}) where T
    valid_insts = collect(1:model.dim.n)
    offset::Int = 0 # Increase as we delete
    for t in axes(model.y, 1), j in axes(model.y, 2)
        if isnan(model.y[t, j])
            deleteat!(valid_insts, t - offset)
            offset += 1
        end
    end
    return valid_insts
end

"""
statespace_likelihood(psitilde::Vector{T}, model::StateSpaceModel, unknowns::Unknowns,
                            valid_insts::Vector{Int}, filter_type::DataType) where T

Compute log-likelihood concerning hyperparameter vector psitilde (``\\psi``)

Evaluate ``\\ell(\\psi;y_n)= -\\frac{np}{2}\\log2\\pi - \\frac{1}{2} \\sum_{t=1}^n \\log |F_t| -
\\frac{1}{2} \\sum_{t=1}^n v_t^{\\top} F_t^{-1} v_t ``
"""
function statespace_likelihood(psitilde::Vector{T}, model::StateSpaceModel, unknowns::Unknowns,
                               valid_insts::Vector{Int}, filter_type::DataType) where T
    # Fill all psitilde in the model
    fill_model_with_parameters!(model, psitilde, unknowns)
    # Calculate v and F
    v, F = get_log_likelihood_params(model, filter_type)
    # Compute log-likelihood based on v and F
    return compute_log_likelihood(model.dim.n, model.dim.p, v, F, valid_insts)
end

"""
    estimate_statespace(model::StateSpaceModel{T}, filter_type::DataType,
                            optimization_method::AbstractOptimizationMethod; verbose::Int = 1) where T

Estimate parameters of the `StateSpaceModel` according to its `filter_type` and `optimization_method`.
"""
function estimate_statespace(model::StateSpaceModel{T}, filter_type::DataType,
                             optimization_method::AbstractOptimizationMethod; verbose::Int = 1) where T
    error(optimization_method , " not implemented") # Returns an error if it cannot
                                                    # find a specialized estimate_statespace
end

function AIC(n_unknowns::Int, log_lik::T) where T
    return T(2 * n_unknowns - 2 * log_lik)
end

function BIC(n::Int, n_unknowns::Int, log_lik::T) where T
    return T(log(n) * n_unknowns - 2 * log_lik)
end

function estimate_statespace!(model::StateSpaceModel{T}, filter_type::DataType,
                              opt_method::AbstractOptimizationMethod{T}; verbose::Int = 1) where T

    unknowns = Unknowns(model)

    if unknowns.n_unknowns == 0
        @warn("StateSpaceModel has no parameters to be estimated.")
        return model
    end

    n_psi          = unknowns.n_unknowns
    loglikelihood = Vector{T}(undef, opt_method.n_seeds)
    psi           = Matrix{T}(undef, n_psi, opt_method.n_seeds)
    optseeds      = Vector{Optim.OptimizationResults}(undef, 0)

    print_estimation_start(verbose, opt_method.n_seeds)

    # compute the valid instants to evaluate loglikelihood
    valid_insts = valid_instants(model)

    t0 = now()

    # Optimization
    for (i, seed) in enumerate(opt_method.seeds)
        try
            optseed = optimize(psitilde -> statespace_likelihood(psitilde, model, unknowns, valid_insts, filter_type), seed,
                            opt_method.method, Optim.Options(f_tol = opt_method.f_tol,
                                                             g_tol = opt_method.g_tol,
                                                             iterations = opt_method.iterations,
                                                             show_trace = (verbose == 3 ? true : false) ))
            loglikelihood[i] = -optseed.minimum
            psi[:, i] = optseed.minimizer
            push!(optseeds, optseed)

            print_loglikelihood(verbose, i, loglikelihood, t0)
        catch err
            println(err)
            println("seed diverged")
        end
    end

    log_lik, best_seed = findmax(loglikelihood)

    aic = AIC(n_psi, log_lik)
    bic = BIC(length(model.y), n_psi, log_lik)

    print_estimation_end(verbose, log_lik, aic, bic)
    verbose >= 2 && println("Best seed optimization result: \n $(optseeds[best_seed])")

    # Take the best seed and fill the model with it
    bestpsi = psi[:, best_seed]
    fill_model_with_parameters!(model, bestpsi, unknowns)

    return
end


# Overloads
export acim, linearresponse, correlationsum, birkhoffcov, birkhoffvar
export lyapunov, covariancefunction, MA_process, AR_process

"""
    acim(L)

Output a Fun object giving the acim of the associated map"""
acim(K::SolutionInvWrapper) = K.op\K.u
acim(L::Operator,u::Fun=uniform(domainspace(L))) = SolutionInv(L,u)\u
acim(M::AbstractMarkovMap) = acim(Transfer(M))

"""
    linearresponse(L, X::Fun)

Output a Fun object giving the first-order change in the acim of the map under perturbation X"""
function linearresponse(K::SolutionInvWrapper,X::Fun)
  @assert all(isapprox.(X.(∂(domain(X))),0;atol=sqrt(eps(cfstype(X))*arclength(domain(X)))))
  K\(-acim(K)*X)'
end

"""
    zero_to(A::Fun, ρ::Fun=uniform(space(A)))

Output ρ*(A - ∫A dρ)"""
function zero_to(A::Fun,ρ::Fun=uniform(space(A)))
  ρA = ρ*A
  ρA -= ρ*sum(ρA)/sum(ρ)
end

"""
    correlationsum(L, A::Fun, ρ=acim(L))

Output resolvent of transfer operator applied to ρ*(A - ∫A dρ)"""
correlationsum(K::SolutionInvWrapper,A::Fun,ρ=acim(K)) = K \ zero_to(A,ρ)

"""
    birkhoffcov(L, A::Fun, B::Fun)

Output covariance of CLT-normalised Birkhoff sums of A and B under the map"""
function birkhoffcov(K::SolutionInvWrapper,A::Fun,B::Fun)
  ρ = acim(K); Az = zero_to(A,ρ); Bz = zero_to(B,ρ)
  sum(B*(K\Az)) + sum(A*(K\Bz)) - sum(B*Az)
end

"""
    birkhoffvar(L, A::Fun)

Output diffusion coefficient of A under the map using Green-Kubo formula"""
function birkhoffvar(K::SolutionInvWrapper,A::Fun)
  Az = zero_to(A,acim(K))
  2sum(A*(K\Az)) - sum(A*Az)
end

for OP in (:linearresponse,:correlationsum,:birkhoffvar)
  @eval $OP(L::Operator,X::Fun) = $OP(SolutionInv(L),X)
  @eval $OP(M::AbstractMarkovMap,X::Fun) = $OP(Transfer(M),X)
  @eval $OP(obj::Operator,X) = $OP(obj,Fun(X,domainspace(obj)))
  @eval $OP(obj::AbstractMarkovMap,X) = $OP(obj,Fun(X,domain(obj)))
end
birkhoffcov(L::Operator,X::Fun,Y::Fun) = birkhoffcov(SolutionInv(L),X,Y)
birkhoffcov(M::AbstractMarkovMap,X::Fun,Y::Fun) = birkhoffcov(Transfer(M),X,Y)

# lag-correlation function
for TYP=(:AbstractMarkovMap,:SolutionInvWrapper)
  @eval covariancefunction(K::$TYP,A::Fun,n::Int) = covariancefunction(Transfer(K),A,n)
  @eval covariancefunction(K::$TYP,A::Fun,B::Fun,n::Int) = covariancefunction(Transfer(K),A,B,n)
  @eval covariancefunction(K::$TYP,A::Fun;tol=eps(norm(A)^2)) = covariancefunction(Transfer(K),A;tol=tol)
  @eval covariancefunction(K::$TYP,A::Fun,B::Fun;tol=eps(norm(A)*norm(B))) = covariancefunction(Transfer(K),A,B;tol=tol)
end
  """
      cfA = covariancefunction(L, A, n)

  Compute the lag-covariance function against transfer operator L for observable A for n time steps
  in each direction. Specifically, the (k+1)th entry is the expectation of A A∘f^(k).
  """
function covariancefunction(L::Operator,A::Fun,n::Int;r=acim(L))
  @compat cf = Array{cfstype(A)}(undef,n+1) # TODO: update to coefficienttype
  Az = zero_to(A,r)
  cf[1] = sum(A*Az)
  for i = 1:n
    Az = L*Az
    Az -= r*sum(Az)
    chop!(Az)
    cf[i+1] = sum(A*Az)
  end
  cf
end

"""
    cf = covariancefunction(L, A; tol = eps(A))

Choose n so that the covariance declines to a given tolerance
"""
function covariancefunction(L::Operator,A::Fun;r=acim(L),tol=eps(norm(A.coefficients)^2))
  n = 16
  @compat cf = Array{cfstype(A)}(undef,n+1)
  Az = zero_to(A,r)
  cf[1] = sum(A*Az)
  for i = 1:n
    Az = L*Az
    Az -= r*sum(Az)
    chop!(Az)
    cf[i+1] = sum(A*Az)
  end
  while maximum(abs.(cf[div(2n,3):end]))>tol && n <= 2^20
    n_old = n
    n = 2n
    pad!(cf,n+1)
    for i = n_old+1:n
      Az = L*Az
      Az -= r*sum(Az)
      chop!(Az)
      cf[i+1] = sum(A*Az)
    end
    # chop!(Az)
  end

  if n==2^20
    warn("Maximum number of coefficients reached (n=$n)")
    cf
  else
    chop!(cf)
  end
end

"""
    cfA, cfB = covariancefunction(L, A, B, n)

Compute the lag-covariance function against transfer operator L between observables A and B for n time steps
in each direction. Specifically, the (k+1)th entry of cfA is the expectation of B A∘f^(k) and vice versa for cfB.
"""
function covariancefunction(L::Operator,A::Fun,B::Fun,n::Int;r=acim(L))
  @compat cfA = Array{cfstype(A)}(undef,n+1)  # TODO: update to coefficienttype
  @compat cfB = Array{cfstype(A)}(undef,n+1)
  Az = zero_to(A,r); Bz = zero_to(B,r)
  cfA[1] = sum(A*Bz)
  cfB[1] = cfA[1]
  for i = 1:n
    Az = L*Az;Bz = L*Bz
    Az -= r*sum(Az); Bz -= r*sum(Bz)
    chop!(Az); chop!(Bz)
    cfA[i+1] = sum(A*Bz); cfB[i+1] = sum(B*Az)
  end
  cfA,cfB
end

"""
    cfA, cfB = covariancefunction(L, A, B; tol=eps(A*B))

Chooses n so that the covariance declines to a given tolerance.
"""
function covariancefunction(L::Operator,A::Fun,B::Fun;r=acim(L),tol=eps(norm(coefficients(A))*norm(coefficients(B))))
  n = 16
  cfA,cfB = covariancefunction(L,A,B,n;r=r)
  while max(maximum(abs.(cfA[div(2n,3):end])),maximum(abs.(cfB[div(2n,3):end])))>tol && n <= 2^20
    n = 2n
    cfA,cfB = covariancefunction(L,A,B,n;r=r)   # chop!(Az)
  end

  if n==2^20
    warn("Maximum number of coefficients reached (n=$n)")
    cfA,cfB
  else
    chop!(cfA),chop!(cfB)
  end
end

# Lyapunov exponent
"""
    lyapunov(f, r=acim(f), sp=Space(rangedomain(f)))

Calculate Lyapunov exponent associated with f"""
function lyapunov(f::AbstractMarkovMap,r=acim(f),sp=Space(rangedomain(f)))
  sum(Fun(x->transferfunction(x,f,x->log(abs(f'(x)))*r(x)),sp))
end
function lyapunov(K::SolutionInvWrapper)
  L = Transfer(K)
  lyapunov(markovmap(L),acim(K),rangespace(L))
end
lyapunov(m::Operator) = lyapunov(markovmap(m),acim(m),rangespace(m))

struct LyapContainer{M<:AbstractMarkovMap,rr}
  m::M
  tr::rr
end
(lc::LyapContainer)(x) = log(abs(lc.m'(x)))*lc.tr(x)

function lyapunov(f::ComposedMarkovMap,r=acim(f))#,sp=Space(rangedomain(f)))
  tr = copy(r)
  lyap = sum(Fun(x->transferfunction(x,f.maps[end],LyapContainer(f.maps[end],tr)),rangedomain(f.maps[end])))
  for k = complength(f)-1:-1:1
    tr = Fun(x->transferfunction(x,f.maps[k+1],tr),rangedomain(f.maps[k+1]))
    lyap += sum(Fun(x->transferfunction(x,f.maps[k],LyapContainer(f.maps[k],tr)),rangedomain(f.maps[k])))
  end
  lyap
end

# Gaussian process parametrisation
function logMA_process(L,A::Fun)
  @assert isreal(A)
  corf = covariancefunction(L,A)
  # d = Circle(); sp = Taylor(d)
  Corfn = Fun(CosSpace(),[corf[1];2corf[2:end]])
  # lCf = log(Corfn)/2
  # SlogMA = Fun(SinSpace(),coefficients(lCf)[2:end])
  # MA = cos(SlogMA)*exp(lCf)
  #  MA = 1+expm1(Fun(sp,chop!(coefficients(log(Corfn)/2))))
  #  # TODO: because the special functions in ApproxFun are buggy
  # chop!(real.(coefficients(MA)[1:2:end]))

  ## OR
  # Corfn = Fun(Laurent(d),[corf[1];repeat(corf[2:end],inner=2)])
  # logMA = Fun(sp,
  # coefficients(Fun(x->log(abs(Corfn(e^(im*x)))),CosSpace())))/2

  logMA = Fun(Taylor(),coefficients(log(Corfn))/2)
end

"""
    MA_process(f,A::Fun)

Calculate coefficients of Gaussian MA process with the same autocorrelation function as A under the map f.

See Appendix A.2 of Wormell, C.L. & Gottwald, G.A. 'On the Validity of Linear Response Theory in High-Dimensional Deterministic Dynamical Systems' (2018).
"""
function MA_process(L,A::Fun)
  logMA = logMA_process(L,A)
  chop!(real.(coefficients(Fun(x->exp(logMA(x)),space(logMA)))))
end

"""
   AR_process(f,A::Fun)

Calculate coefficients of Gaussian AR process with the same autocorrelation function as A under the map f.

See Appendix A.2 of Wormell, C.L. & Gottwald, G.A. 'On the Validity of Linear Response Theory in High-Dimensional Deterministic Dynamical Systems' (2018).
"""
function AR_process(L,A::Fun)
  logMA = logMA_process(L,A)
  chop!(real.(coefficients(Fun(x->exp(-logMA(x)),space(logMA)))))
end


using LinearTransforms, Arpack
using SparseArrays, LinearAlgebra
using Test

@testset "opMask" begin

    m, n = 123, 321
    M = sprand(Bool,m,n,0.5)
    I, J = findnz(M)[1:2]
    u, v = randn(m,2), randn(n,2)
    X = u*v'

    @testset "full matrices" begin
        op = opMask(M)

        # forward operator
        w = op*X
        wt = [X[i,j] for (i,j) in zip(I,J)]
        @test norm(w-wt, Inf) ≤ 1e-10*(1+norm(wt,Inf))

        # adjoint operator
        V = op'*w
        Vt = sparse(I, J, w, m, n)
        @test norm(V-Vt) ≤ 1e-10

        # dot test
        @test dot(op*X,w) ≈ dot(X,op'*w)

        # test eltypes
        op = opMask(M)
        @test eltype(op) == Float64
        op = opMask{Complex{Float64}}(M)
        @test eltype(op) == Complex{Float64}

        # test size
        @test size(op) == (nnz(M), size(M))

    end

    @testset "factored matrices" begin
        op = opMaskFactored(M)

        # forward operator
        w = op*(u,v)
        wT = [X[i,j] for (i,j) in zip(I,J)]
        @test norm(w-wT, Inf) ≤ 1e-10*(1+norm(wT,Inf))

        # adjoint operator
        V = op'*w
        Vt = sparse(I, J, w, m, n)
        @test norm(V-Vt) ≤ 1e-10

        # dot test
        @test dot(op*(u,v),w) ≈ dot(u*v',op'*w)

        # eltypes
        op = opMaskFactored(M)
        @test eltype(op) == Float64
        op = opMaskFactored{Int64}(M)
        @test eltype(op) == Int64

        # test size
        @test size(op) == (nnz(M), size(M))

    end
end


"""
Tests the forward modelling function `calculate_response` against a range of
models. For model information, see "./test-models/README.md". A comparison
of the expected output and calculated output for each model is plotted and saved
into "./test-models/generated-plots".
"""
module TestCalcResponse
using Base.Test
using MT1D
using Utils
using Gadfly

fs = logspace(-4, 4, 100)
cd("./test-models") do
    model_files = filter(s -> ismatch(r"\.model$", s), readdir())

    map(model_files) do fname
        model = readdlm(fname)
        expected_response = readdlm("$fname.response")

        ρ, Φ = calc_response(model, fs)
        p = plot_response(expected_response, [ρ Φ], model, fs)
        draw(SVG("./generated-plots/$fname.svg", 1140px, 800px), p)

        @test_approx_eq expected_response[:,1] ρ
        @test_approx_eq expected_response[:,2] Φ
    end
end

end