import Base.LinAlg.BLAS
using CUBLAS
using CUDArt
using Base.Test

m = 20
n = 35
k = 13

function blasabs(A)
    return abs.(real(A)) + abs.(imag(A))
end

#################
# level 1 tests #
#################

@testset "blascopy!" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        A = convert(Vector{elty}, collect(1:m))
        if elty <: Base.LinAlg.BlasComplex
            A += im*A
        end
        @test ndims(A) == 1
        n1 = length(A)
        d_A = CudaArray(A)
        d_B = CudaArray(elty, n1)
        CUBLAS.blascopy!(n,d_A,1,d_B,1)
        B = to_host(d_B)
        @test A == B
    end
end

@testset "scal!" begin
    function test_scal!{T}(alpha,A::Array{T})
        @test ndims(A) == 1
        n1 = length(A)
        d_A = CudaArray(A)
        CUBLAS.scal!(n1,alpha,d_A,1)
        A1 = to_host(d_A)
        @test alpha*A ≈ A1

        d_A = CudaArray(A)
        d_As = CUBLAS.scale(d_A, alpha)
        A1 = to_host(d_As)
        @test alpha*A ≈ A1

        CUBLAS.scale!(d_A, alpha)
        A1 = to_host(d_As)
        @test alpha*A ≈ A1
    end
    test_scal!(2.0f0,Float32[1:m;])
    test_scal!(2.0,Float64[1:m;])
    test_scal!(1.0f0+im*1.0f0,Float32[1:m;]+im*Float32[1:m;])
    test_scal!(1.0+im*1.0,Float64[1:m;]+im*Float64[1:m;])
    test_scal!(2.0f0,Float32[1:m;]+im*Float32[1:m;])
    test_scal!(2.0,Float64[1:m;]+im*Float64[1:m;])
end

@testset "dot" begin
    @testset for elty in [Float32, Float64]
        A = convert(Vector{elty}, collect(1:m))
        B = convert(Vector{elty}, collect(1:m))
        @test ndims(A) == 1
        @test ndims(B) == 1
        @test length(A) == length(B)
        n1 = length(A)
        d_A = CudaArray(A)
        d_B = CudaArray(B)
        cuda_dot1 = CUBLAS.dot(n1,d_A,1,d_B,1)
        cuda_dot2 = CUBLAS.dot(d_A,d_B)
        host_dot = dot(A,B)
        @test host_dot ≈ cuda_dot1
        @test host_dot ≈ cuda_dot2

        #d_A = CudaArray(A)
        #d_B = CudaArray(B)
        #cuda_dot3 = CUBLAS.dot(d_A, 3:5, d_B, 5:7)
        #host_dot3 = dot(A, 3:5, B, 5:7)
        #@test_approx_eq(cuda_dot3, host_dot3)
    end
end

@testset "dotu" begin
    @testset for elty in [Complex64, Complex128]
        A = rand(elty, m)
        B = rand(elty, m)
        @test ndims(A) == 1
        @test ndims(B) == 1
        @test length(A) == length(B)
        n1 = length(A)
        d_A = CudaArray(A)
        d_B = CudaArray(B)
        cuda_dot1 = CUBLAS.dotu(n1,d_A,1,d_B,1)
        cuda_dot2 = CUBLAS.dotu(d_A,d_B)
        host_dot = A.'*B
        if VERSION < v"0.6.0-dev.2074" # julia PR #19670
            @test host_dot[1] ≈ cuda_dot1
            @test host_dot[1] ≈ cuda_dot2
            @test host_dot ≈ (d_A.'*d_B)
        else
            @test host_dot ≈ cuda_dot1
            @test host_dot ≈ cuda_dot2
            @test host_dot ≈ (d_A.'*d_B)[1]
        end
    end
end

@testset "dotc" begin
    @testset for elty in [Complex64, Complex128]
        A = rand(elty, m)
        B = rand(elty, m)
        @test ndims(A) == 1
        @test ndims(B) == 1
        @test length(A) == length(B)
        n1 = length(A)
        d_A = CudaArray(A)
        d_B = CudaArray(B)
        cuda_dot1 = CUBLAS.dotc(n1,d_A,1,d_B,1)
        cuda_dot2 = CUBLAS.dotc(d_A,d_B)
        host_dot = A'*B
        if VERSION < v"0.6.0-dev.2074" # julia PR #19670
            @test host_dot[1] ≈ cuda_dot1
            @test host_dot[1] ≈ cuda_dot2
            @test host_dot ≈ (d_A'*d_B)
        else
            @test host_dot ≈ cuda_dot1
            @test host_dot ≈ cuda_dot2
            @test host_dot ≈ (d_A'*d_B)[1]
        end
    end
end

@testset "nrm2" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        A = rand(elty, m)
        @test ndims(A) == 1
        n1 = length(A)
        d_A = CudaArray(A)
        cuda_nrm2_1 = CUBLAS.nrm2(n1,d_A,1)
        cuda_nrm2_2 = CUBLAS.nrm2(d_A)
        cuda_nrm2_3 = norm(d_A)
        host_nrm2 = norm(A)
        @test host_nrm2 ≈ cuda_nrm2_1
        @test host_nrm2 ≈ cuda_nrm2_2
        @test host_nrm2 ≈ cuda_nrm2_3
    end
end

@testset "asum" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        A = rand(elty, m)
        @test ndims(A) == 1
        n1 = length(A)
        d_A = CudaArray(A)
        cuda_asum1 = CUBLAS.asum(n1,d_A,1)
        cuda_asum2 = CUBLAS.asum(d_A)
        host_asum = sum(abs.(real(A)) + abs.(imag(A)))
        @test host_asum ≈ cuda_asum1
        @test host_asum ≈ cuda_asum2
    end
end

@testset "axpy!" begin
    # test axpy!
    function test_axpy!_1(alpha,A,B)
        @test length(A) == length(B)
        n1 = length(A)
        d_A = CudaArray(A)
        d_B1 = CudaArray(B)
        CUBLAS.axpy!(n1,alpha,d_A,1,d_B1,1)
        B1 = to_host(d_B1)
        host_axpy = alpha*A + B
        @test host_axpy ≈ B1
    end
    test_axpy!_1(2.0f0,rand(Float32,m),rand(Float32,m))
    test_axpy!_1(2.0,rand(Float64,m),rand(Float64,m))
    test_axpy!_1(2.0f0+im*2.0f0,rand(Complex64,m),rand(Complex64,m))
    test_axpy!_1(2.0+im*2.0,rand(Complex128,m),rand(Complex128,m))

    function test_axpy!_2(alpha,A,B)
        @test length(A) == length(B)
        n1 = length(A)
        d_A = CudaArray(A)
        d_B1 = CudaArray(B)
        CUBLAS.axpy!(alpha,d_A,d_B1)
        B1 = to_host(d_B1)
        host_axpy = alpha*A + B
        @test host_axpy ≈ B1
    end
    test_axpy!_2(2.0f0,rand(Float32,m),rand(Float32,m))
    test_axpy!_2(2.0,rand(Float64,m),rand(Float64,m))
    test_axpy!_2(2.0f0+im*2.0f0,rand(Complex64,m),rand(Complex64,m))
    test_axpy!_2(2.0+im*2.0,rand(Complex128,m),rand(Complex128,m))

    #=function test_axpy!_3(alpha,A,B)
        @test length(A) == length(B)
        n1 = length(A)
        d_A = CudaArray(A)
        d_B1 = CudaArray(B)
        CUBLAS.axpy!(alpha,d_A,1:2:n1,d_B1,1:2:n1)
        B1 = to_host(d_B1)
        host_axpy = B
        host_axpy[1:2:n1] = alpha*A[1:2:n1] + B[1:2:n1]
        @test_approx_eq(host_axpy,B1)
    end
    test_axpy!_3(2.0f0,rand(Float32,m),rand(Float32,m))
    test_axpy!_3(2.0,rand(Float64,m),rand(Float64,m))
    test_axpy!_3(2.0f0+im*2.0f0,rand(Complex64,m),rand(Complex64,m))
    test_axpy!_3(2.0+im*2.0,rand(Complex128,m),rand(Complex128,m))

    function test_axpy!_4(alpha,A,B)
        @test length(A) == length(B)
        n1 = length(A)
        d_A = CudaArray(A)
        d_B1 = CudaArray(B)
        r = 1:div(n1,2)
        CUBLAS.axpy!(alpha,d_A,r,d_B1,r)
        B1 = to_host(d_B1)
        host_axpy = B
        host_axpy[r] = alpha*A[r] + B[r]
        @test_approx_eq(host_axpy,B1)
    end
    test_axpy!_4(2.0f0,rand(Float32,m),rand(Float32,m))
    test_axpy!_4(2.0,rand(Float64,m),rand(Float64,m))
    test_axpy!_4(2.0f0+im*2.0f0,rand(Complex64,m),rand(Complex64,m))
    test_axpy!_4(2.0+im*2.0,rand(Complex128,m),rand(Complex128,m))=#
end

@testset "iamax and iamin" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        A = rand(elty, m)
        n1 = length(A)
        d_A = CudaArray(A)
        Aabs = blasabs(A)
        imin1 = CUBLAS.iamin(n1,d_A,1)
        imax1 = CUBLAS.iamax(n1,d_A,1)
        imin2 = CUBLAS.iamin(d_A)
        imax2 = CUBLAS.iamax(d_A)
        host_imin = indmin(Aabs)
        host_imax = indmax(Aabs)
        @test imin1 == imin2 == host_imin
        @test imin1 == imin2 == host_imin
    end
end

#################
# level 2 tests #
#################

@testset "gemv!" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        alpha = convert(elty,1)
        beta = convert(elty,1)
        A = rand(elty,m,n)
        d_A = CudaArray(A)

        # test y = A*x + y
        x = rand(elty,n)
        d_x = CudaArray(x)
        y = rand(elty,m)
        d_y = CudaArray(y)
        y = A*x + y
        CUBLAS.gemv!('N',alpha,d_A,d_x,beta,d_y)
        h_y = to_host(d_y)
        @test y ≈ h_y
        A_mul_B!(d_y,d_A,d_x)
        h_y = to_host(d_y)
        @test h_y ≈ A*x

        # test x = A.'*y + x
        x = rand(elty,n)
        d_x = CudaArray(x)
        y = rand(elty,m)
        d_y = CudaArray(y)
        x = A.'*y + x
        CUBLAS.gemv!('T',alpha,d_A,d_y,beta,d_x)
        h_x = to_host(d_x)
        @test x ≈ h_x
        At_mul_B!(d_x,d_A,d_y)
        h_x = to_host(d_x)
        @test h_x ≈ A.'*y

        # test x = A'*y + x
        x = rand(elty,n)
        d_x = CudaArray(x)
        y = rand(elty,m)
        d_y = CudaArray(y)
        x = A'*y + x
        CUBLAS.gemv!('C',alpha,d_A,d_y,beta,d_x)
        h_x = to_host(d_x)
        @test x ≈ h_x
        Ac_mul_B!(d_x,d_A,d_y)
        h_x = to_host(d_x)
        @test h_x ≈ A'*y
    end
end

@testset "gemv" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        alpha = convert(elty,2)
        A = rand(elty,m,n)
        d_A = CudaArray(A)
        # test y = alpha*(A*x)
        x = rand(elty,n)
        d_x = CudaArray(x)
        y1 = alpha*(A*x)
        y2 = A*x
        d_y1 = CUBLAS.gemv('N',alpha,d_A,d_x)
        d_y2 = CUBLAS.gemv('N',d_A,d_x)
        h_y1 = to_host(d_y1)
        h_y2 = to_host(d_y2)
        @test y1 ≈ h_y1
        @test y2 ≈ h_y2
        @test y2 ≈ to_host(d_A * d_x)

        # test x = alpha*(A.'*y)
        y = rand(elty,m)
        d_y = CudaArray(y)
        x1 = alpha*(A.'*y)
        x2 = A.'*y
        d_x1 = CUBLAS.gemv('T',alpha,d_A,d_y)
        d_x2 = CUBLAS.gemv('T',d_A,d_y)
        h_x1 = to_host(d_x1)
        h_x2 = to_host(d_x2)
        @test x1 ≈ h_x1
        @test x2 ≈ h_x2
        @test x2 ≈ to_host(d_A.' * d_y)

        # test x = alpha*(A'*y)
        y = rand(elty,m)
        d_y = CudaArray(y)
        x1 = alpha*(A'*y)
        x2 = A'*y
        d_x1 = CUBLAS.gemv('C',alpha,d_A,d_y)
        d_x2 = CUBLAS.gemv('C',d_A,d_y)
        h_x1 = to_host(d_x1)
        h_x2 = to_host(d_x2)
        @test y1 ≈ h_y1
        @test y2 ≈ h_y2
        @test x2 ≈ to_host(d_A' * d_y)
    end
end

@testset "gbmv!" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # parameters
        alpha = convert(elty,2)
        beta = convert(elty,3)
        # bands
        ku = 2
        kl = 3
        # generate banded matrix
        A = rand(elty,m,n)
        A = bandex(A,kl,ku)
        # get packed format
        Ab = band(A,kl,ku)
        d_Ab = CudaArray(Ab)
        # test y = alpha*A*x + beta*y
        x = rand(elty,n)
        d_x = CudaArray(x)
        y = rand(elty,m)
        d_y = CudaArray(y)
        CUBLAS.gbmv!('N',m,kl,ku,alpha,d_Ab,d_x,beta,d_y)
        BLAS.gbmv!('N',m,kl,ku,alpha,Ab,x,beta,y)
        h_y = to_host(d_y)
        @test y ≈ h_y
        # test y = alpha*A.'*x + beta*y
        x = rand(elty,n)
        d_x = CudaArray(x)
        y = rand(elty,m)
        d_y = CudaArray(y)
        CUBLAS.gbmv!('T',m,kl,ku,alpha,d_Ab,d_y,beta,d_x)
        BLAS.gbmv!('T',m,kl,ku,alpha,Ab,y,beta,x)
        h_x = to_host(d_x)
        @test x ≈ h_x
        # test y = alpha*A'*x + beta*y
        x = rand(elty,n)
        d_x = CudaArray(x)
        y = rand(elty,m)
        d_y = CudaArray(y)
        CUBLAS.gbmv!('C',m,kl,ku,alpha,d_Ab,d_y,beta,d_x)
        BLAS.gbmv!('C',m,kl,ku,alpha,Ab,y,beta,x)
        h_x = to_host(d_x)
        @test x ≈ h_x
    end
end

@testset "gbmv" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # parameters
        alpha = convert(elty,2)
        # bands
        ku = 2
        kl = 3
        # generate banded matrix
        A = rand(elty,m,n)
        A = bandex(A,kl,ku)
        # get packed format
        Ab = band(A,kl,ku)
        d_Ab = CudaArray(Ab)
        # test y = alpha*A*x
        x = rand(elty,n)
        d_x = CudaArray(x)
        d_y = CUBLAS.gbmv('N',m,kl,ku,alpha,d_Ab,d_x)
        y = zeros(elty,m)
        y = BLAS.gbmv('N',m,kl,ku,alpha,Ab,x)
        h_y = to_host(d_y)
        @test y ≈ h_y
    end
end

@testset "symv!" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # parameters
        alpha = convert(elty,2)
        beta = convert(elty,3)
        # generate symmetric matrix
        A = rand(elty,m,m)
        A = A + A.'
        # generate vectors
        x = rand(elty,m)
        y = rand(elty,m)
        # copy to device
        d_A = CudaArray(A)
        d_x = CudaArray(x)
        d_y = CudaArray(y)
        # execute on host
        BLAS.symv!('U',alpha,A,x,beta,y)
        # execute on device
        CUBLAS.symv!('U',alpha,d_A,d_x,beta,d_y)
        # compare results
        h_y = to_host(d_y)
        @test y ≈ h_y
    end
end

@testset "symv" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # generate symmetric matrix
        A = rand(elty,m,m)
        A = A + A.'
        # generate vectors
        x = rand(elty,m)
        # copy to device
        d_A = CudaArray(A)
        d_x = CudaArray(x)
        # execute on host
        y = BLAS.symv('U',A,x)
        # execute on device
        d_y = CUBLAS.symv('U',d_A,d_x)
        # compare results
        h_y = to_host(d_y)
        @test y ≈ h_y
    end
end

@testset "hemv!" begin
    @testset for elty in [Complex64, Complex128]
        # parameters
        alpha = convert(elty,2)
        beta = convert(elty,3)
        # generate hermitian matrix
        A = rand(elty,m,m)
        A = A + A'
        # generate vectors
        x = rand(elty,m)
        y = rand(elty,m)
        # copy to device
        d_A = CudaArray(A)
        d_x = CudaArray(x)
        d_y = CudaArray(y)
        # execute on host
        BLAS.hemv!('U',alpha,A,x,beta,y)
        # execute on device
        CUBLAS.hemv!('U',alpha,d_A,d_x,beta,d_y)
        # compare results
        h_y = to_host(d_y)
        @test y ≈ h_y
    end
end

@testset "hemv" begin
    @testset for elty in [Complex64, Complex128]
        # generate hermitian matrix
        A = rand(elty,m,m)
        A = A + A.'
        # generate vectors
        x = rand(elty,m)
        # copy to device
        d_A = CudaArray(A)
        d_x = CudaArray(x)
        # execute on host
        y = BLAS.hemv('U',A,x)
        # execute on device
        d_y = CUBLAS.hemv('U',d_A,d_x)
        # compare results
        h_y = to_host(d_y)
        @test y ≈ h_y
    end
end

@testset "sbmv!" begin
    @testset for elty in [Float32, Float64]
        # parameters
        alpha = convert(elty,3)
        beta = convert(elty,2.5)
        # generate symmetric matrix
        A = rand(elty,m,m)
        A = A + A'
        # restrict to 3 bands
        nbands = 3
        @test m >= 1+nbands
        A = bandex(A,nbands,nbands)
        # convert to 'upper' banded storage format
        AB = band(A,0,nbands)
        # construct x and y
        x = rand(elty,m)
        y = rand(elty,m)
        # move to host
        d_AB = CudaArray(AB)
        d_x = CudaArray(x)
        d_y = CudaArray(y)
        # sbmv!
        CUBLAS.sbmv!('U',nbands,alpha,d_AB,d_x,beta,d_y)
        y = alpha*(A*x) + beta*y
        # compare
        h_y = to_host(d_y)
        @test y ≈ h_y
    end
end

@testset "sbmv" begin
    @testset for elty in [Float32, Float64]
        # parameters
        alpha = convert(elty,3)
        beta = convert(elty,2.5)
        # generate symmetric matrix
        A = rand(elty,m,m)
        A = A + A'
        # restrict to 3 bands
        nbands = 3
        @test m >= 1+nbands
        A = bandex(A,nbands,nbands)
        # convert to 'upper' banded storage format
        AB = band(A,0,nbands)
        # construct x and y
        x = rand(elty,m)
        y = rand(elty,m)
        # move to host
        d_AB = CudaArray(AB)
        d_x = CudaArray(x)
        # sbmv!
        d_y = CUBLAS.sbmv('U',nbands,d_AB,d_x)
        y = A*x
        # compare
        h_y = to_host(d_y)
        @test y ≈ h_y
    end
end

@testset "hbmv!" begin
    @testset for elty in [Complex64, Complex128]
        # parameters
        alpha = rand(elty)
        beta = rand(elty)
        # generate Hermitian matrix
        A = rand(elty,m,m)
        A = A + ctranspose(A)
        # restrict to 3 bands
        nbands = 3
        @test m >= 1+nbands
        A = bandex(A,nbands,nbands)
        # convert to 'upper' banded storage format
        AB = band(A,0,nbands)
        # construct x and y
        x = rand(elty,m)
        y = rand(elty,m)
        # move to host
        d_AB = CudaArray(AB)
        d_x = CudaArray(x)
        d_y = CudaArray(y)
        # hbmv!
        CUBLAS.hbmv!('U',nbands,alpha,d_AB,d_x,beta,d_y)
        y = alpha*(A*x) + beta*y
        # compare
        h_y = to_host(d_y)
        @test y ≈ h_y
    end
end

@testset "hbmv" begin
    @testset for elty in [Complex64, Complex128]
        # parameters
        alpha = rand(elty)
        beta = rand(elty)
        # generate Hermitian matrix
        A = rand(elty,m,m)
        A = A + ctranspose(A)
        # restrict to 3 bands
        nbands = 3
        @test m >= 1+nbands
        A = bandex(A,nbands,nbands)
        # convert to 'upper' banded storage format
        AB = band(A,0,nbands)
        # construct x and y
        x = rand(elty,m)
        y = rand(elty,m)
        # move to host
        d_AB = CudaArray(AB)
        d_x = CudaArray(x)
        # hbmv
        d_y = CUBLAS.hbmv('U',nbands,d_AB,d_x)
        y = A*x
        # compare
        h_y = to_host(d_y)
        @test y ≈ h_y
    end
end

@testset "tbmv!" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # generate triangular matrix
        A = rand(elty,m,m)
        # restrict to 3 bands
        nbands = 3
        @test m >= 1+nbands
        A = bandex(A,0,nbands)
        # convert to 'upper' banded storage format
        AB = band(A,0,nbands)
        # construct x and y
        x = rand(elty,m)
        # move to host
        d_AB = CudaArray(AB)
        d_x = CudaArray(x)
        # tbmv!
        CUBLAS.tbmv!('U','N','N',nbands,d_AB,d_x)
        x = A*x
        # compare
        h_x = to_host(d_x)
        @test x ≈ h_x
    end
end

@testset "tbmv" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # generate triangular matrix
        A = rand(elty,m,m)
        # restrict to 3 bands
        nbands = 3
        @test m >= 1+nbands
        A = bandex(A,0,nbands)
        # convert to 'upper' banded storage format
        AB = band(A,0,nbands)
        # construct x
        x = rand(elty,m)
        # move to host
        d_AB = CudaArray(AB)
        d_x = CudaArray(x)
        # tbmv!
        d_y = CUBLAS.tbmv!('U','N','N',nbands,d_AB,d_x)
        y = A*x
        # compare
        h_y = to_host(d_y)
        @test y ≈ h_y
    end
end

@testset "tbsv!" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # generate triangular matrix
        A = rand(elty,m,m)
        # restrict to 3 bands
        nbands = 3
        @test m >= 1+nbands
        A = bandex(A,0,nbands)
        # convert to 'upper' banded storage format
        AB = band(A,0,nbands)
        # generate vector
        x = rand(elty,m)
        # move to device
        d_AB = CudaArray(AB)
        d_x = CudaArray(x)
        #tbsv!
        CUBLAS.tbsv!('U','N','N',nbands,d_AB,d_x)
        x = A\x
        # compare
        h_x = to_host(d_x)
        @test x ≈ h_x
    end
end

@testset "tbsv" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # generate triangular matrix
        A = rand(elty,m,m)
        # restrict to 3 bands
        nbands = 3
        @test m >= 1+nbands
        A = bandex(A,0,nbands)
        # convert to 'upper' banded storage format
        AB = band(A,0,nbands)
        # generate vector
        x = rand(elty,m)
        # move to device
        d_AB = CudaArray(AB)
        d_x = CudaArray(x)
        #tbsv
        d_y = CUBLAS.tbsv('U','N','N',nbands,d_AB,d_x)
        y = A\x
        # compare
        h_y = to_host(d_y)
        @test y ≈ h_y
    end
end

@testset "trmv!" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # generate triangular matrix
        A = rand(elty,m,m)
        A = triu(A)
        # generate vector
        x = rand(elty,m)
        # move to device
        d_A = CudaArray(A)
        d_x = CudaArray(x)
        # execute trmv!
        CUBLAS.trmv!('U','N','N',d_A,d_x)
        x = A*x
        # compare
        h_x = to_host(d_x)
        @test x ≈ h_x
    end
end

@testset "trmv" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # generate triangular matrix
        A = rand(elty,m,m)
        A = triu(A)
        # generate vector
        x = rand(elty,m)
        # move to device
        d_A = CudaArray(A)
        d_x = CudaArray(x)
        # execute trmv!
        d_y = CUBLAS.trmv('U','N','N',d_A,d_x)
        y = A*x
        # compare
        h_y = to_host(d_y)
        @test y ≈ h_y
    end
end

@testset "trsv!" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # generate triangular matrix
        A = rand(elty,m,m)
        A = triu(A)
        # generate vector
        x = rand(elty,m)
        # move to device
        d_A = CudaArray(A)
        d_x = CudaArray(x)
        # execute trsv!
        CUBLAS.trsv!('U','N','N',d_A,d_x)
        x = A\x
        # compare
        h_x = to_host(d_x)
        @test x ≈ h_x
    end
end

@testset "trsv" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # generate triangular matrix
        A = rand(elty,m,m)
        A = triu(A)
        # generate vector
        x = rand(elty,m)
        # move to device
        d_A = CudaArray(A)
        d_x = CudaArray(x)
        # execute trsv!
        d_y = CUBLAS.trsv('U','N','N',d_A,d_x)
        y = A\x
        # compare
        h_y = to_host(d_y)
        @test y ≈ h_y
    end
end

@testset "ger!" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # construct matrix and vectors
        A = rand(elty,m,n)
        x = rand(elty,m)
        y = rand(elty,n)
        alpha = convert(elty,2)
        # move to device
        d_A = CudaArray(A)
        d_x = CudaArray(x)
        d_y = CudaArray(y)
        # perform rank one update
        CUBLAS.ger!(alpha,d_x,d_y,d_A)
        A = (alpha*x)*y' + A
        # move to host and compare
        h_A = to_host(d_A)
        @test A ≈ h_A
    end
end

@testset "syr!" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # construct matrix and vector
        A = rand(elty,m,m)
        A = A + A.'
        x = rand(elty,m)
        alpha = convert(elty,2)
        # move to device
        d_A = CudaArray(A)
        d_x = CudaArray(x)
        # perform rank one update
        CUBLAS.syr!('U',alpha,d_x,d_A)
        A = (alpha*x)*x.' + A
        # move to host and compare upper triangles
        h_A = to_host(d_A)
        A = triu(A)
        h_A = triu(h_A)
        @test A ≈ h_A
    end
end

@testset "her!" begin
    @testset for elty in [Complex64, Complex128]
        local m = 2
        # construct matrix and vector
        A = rand(elty,m,m)
        A = A + A'
        x = rand(elty,m)
        alpha = convert(elty,2)
        # move to device
        d_A = CudaArray(A)
        d_x = CudaArray(x)
        # perform rank one update
        CUBLAS.her!('U',alpha,d_x,d_A)
        A = (alpha*x)*x' + A
        # move to host and compare upper triangles
        h_A = to_host(d_A)
        A = triu(A)
        h_A = triu(h_A)
        @test A ≈ h_A
    end
end

@testset "her2!" begin
    @testset for elty in [Complex64, Complex128]
        local m = 2
        # construct matrix and vector
        A = rand(elty,m,m)
        A = A + A'
        x = rand(elty,m)
        y = rand(elty,m)
        alpha = convert(elty,2)
        # move to device
        d_A = CudaArray(A)
        d_x = CudaArray(x)
        d_y = CudaArray(y)
        # perform rank one update
        CUBLAS.her2!('U',alpha,d_x,d_y,d_A)
        A = (alpha*x)*y' + y*(alpha*x)' + A
        # move to host and compare upper triangles
        h_A = to_host(d_A)
        A = triu(A)
        h_A = triu(h_A)
        @test A ≈ h_A
    end
end

@testset "gemm!" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # parameters
        alpha = rand(elty)
        beta = rand(elty)
        # generate matrices
        A = rand(elty,m,k)
        B = rand(elty,k,n)
        C1 = rand(elty,m,n)
        C2 = copy(C1)
        # move to device
        d_A = CudaArray(A)
        d_B = CudaArray(B)
        d_C1 = CudaArray(C1)
        d_C2 = CudaArray(C2)
        # C = (alpha*A)*B + beta*C
        CUBLAS.gemm!('N','N',alpha,d_A,d_B,beta,d_C1)
        A_mul_B!(d_C2, d_A, d_B)
        h_C1 = to_host(d_C1)
        h_C2 = to_host(d_C2)
        C1 = (alpha*A)*B + beta*C1
        C2 = A*B
        # compare
        @test C1 ≈ h_C1
        @test C2 ≈ h_C2
    end
end

@testset "gemm" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # generate matrices
        A = rand(elty,m,k)
        B = rand(elty,k,n)
        # move to device
        d_A = CudaArray(A)
        d_B = CudaArray(B)
        # C = (alpha*A)*B + beta*C
        d_C = CUBLAS.gemm('N','N',d_A,d_B)
        C = A*B
        C2 = d_A * d_B
        # compare
        h_C = to_host(d_C)
        h_C2 = to_host(C2)
        @test C ≈ h_C
        @test C ≈ h_C2
    end
end

@testset "gemm_batched!" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # parameters
        alpha = rand(elty)
        beta = rand(elty)
        # generate matrices
        A = [rand(elty,m,k) for i in 1:10]
        B = [rand(elty,k,n) for i in 1:10]
        C = [rand(elty,m,n) for i in 1:10]
        # move to device
        d_A = CudaArray{elty, 2}[]
        d_B = CudaArray{elty, 2}[]
        d_C = CudaArray{elty, 2}[]
        for i in 1:length(A)
            push!(d_A,CudaArray(A[i]))
            push!(d_B,CudaArray(B[i]))
            push!(d_C,CudaArray(C[i]))
        end
        # C = (alpha*A)*B + beta*C
        CUBLAS.gemm_batched!('N','N',alpha,d_A,d_B,beta,d_C)
        for i in 1:length(d_C)
            C[i] = (alpha*A[i])*B[i] + beta*C[i]
            h_C = to_host(d_C[i])
            #compare
            @test C[i] ≈ h_C
        end
    end
end

@testset "gemm_batched" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # generate matrices
        A = [rand(elty,m,k) for i in 1:10]
        B = [rand(elty,k,n) for i in 1:10]
        # move to device
        d_A = CudaArray{elty, 2}[]
        d_B = CudaArray{elty, 2}[]
        for i in 1:length(A)
            push!(d_A, CudaArray(A[i]))
            push!(d_B, CudaArray(B[i]))
        end
        # C = A*B
        d_C = CUBLAS.gemm_batched('N','N',d_A,d_B)
        for i in 1:length(A)
            C = A[i]*B[i]
            h_C = to_host(d_C[i])
            @test C ≈ h_C
        end
    end
end

@testset "symm!" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # parameters
        alpha = rand(elty)
        beta = rand(elty)
        # generate matrices
        A = rand(elty,m,m)
        A = A + A.'
        B = rand(elty,m,n)
        C = rand(elty,m,n)
        # move to device
        d_A = CudaArray(A)
        d_B = CudaArray(B)
        d_C = CudaArray(C)
        # C = (alpha*A)*B + beta*C
        CUBLAS.symm!('L','U',alpha,d_A,d_B,beta,d_C)
        C = (alpha*A)*B + beta*C
        # compare
        h_C = to_host(d_C)
        @test C ≈ h_C
    end
end

@testset "symm" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # generate matrices
        A = rand(elty,m,m)
        A = A + A.'
        B = rand(elty,m,n)
        # move to device
        d_A = CudaArray(A)
        d_B = CudaArray(B)
        # C = (alpha*A)*B + beta*C
        d_C = CUBLAS.symm('L','U',d_A,d_B)
        C = A*B
        # compare
        h_C = to_host(d_C)
        @test C ≈ h_C
    end
end

@testset "syrk!" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # generate matrices
        A = rand(elty,m,k)
        C = rand(elty,m,m)
        C = C + C.'
        # parameters
        alpha = rand(elty)
        beta = rand(elty)
        # move to device
        d_A = CudaArray(A)
        d_C = CudaArray(C)
        # C = (alpha*A)*A.' + beta*C
        CUBLAS.syrk!('U','N',alpha,d_A,beta,d_C)
        C = (alpha*A)*A.' + beta*C
        C = triu(C)
        # move to host and compare
        h_C = to_host(d_C)
        h_C = triu(C)
        @test C ≈ h_C
    end
end

@testset "syrk" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # generate matrices
        A = rand(elty,m,k)
        # move to device
        d_A = CudaArray(A)
        # C = A*A.'
        d_C = CUBLAS.syrk('U','N',d_A)
        C = A*A.'
        C = triu(C)
        # move to host and compare
        h_C = to_host(d_C)
        h_C = triu(C)
        @test C ≈ h_C
    end
end

@testset "herk!" begin
    @testset for elty in [Complex64, Complex128]
        # generate matrices
        A = rand(elty,m,k)
        C = rand(elty,m,m)
        C = C + C'
        # parameters
        alpha = rand(elty)
        beta = rand(elty)
        # move to device
        d_A = CudaArray(A)
        d_C = CudaArray(C)
        CUBLAS.herk!('U','N',alpha,d_A,beta,d_C)
        C = alpha*(A*A') + beta*C
        C = triu(C)
        # move to host and compare
        h_C = to_host(d_C)
        h_C = triu(C)
        @test C ≈ h_C
    end
end

@testset "herk" begin
    @testset for elty in [Complex64, Complex128]
        # generate matrices
        A = rand(elty,m,k)
        # move to device
        d_A = CudaArray(A)
        # C = A*A'
        d_C = CUBLAS.herk('U','N',d_A)
        C = A*A'
        C = triu(C)
        # move to host and compare
        h_C = to_host(d_C)
        h_C = triu(C)
        @test C ≈ h_C
    end
end

@testset "syr2k!" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        #local m = 3
        #local k = 1
        # generate parameters
        alpha = rand(elty)
        beta = rand(elty)
        # generate matrices
        A = rand(elty,m,k)
        B = rand(elty,m,k)
        C = rand(elty,m,m)
        C = C + C.'
        # move to device
        d_A = CudaArray(A)
        d_B = CudaArray(B)
        d_C = CudaArray(C)
        # compute
        #C = alpha*(A*B.') + conj(alpha)*(B*A.') + beta*C
        C = alpha*(A*B.' + B*A.') + beta*C
        CUBLAS.syr2k!('U','N',alpha,d_A,d_B,beta,d_C)
        # move back to host and compare
        C = triu(C)
        h_C = to_host(d_C)
        h_C = triu(h_C)

        @test C ≈ h_C
    end
end

@testset "syr2k" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # generate parameters
        alpha = rand(elty)
        # generate matrices
        A = rand(elty,m,k)
        B = rand(elty,m,k)
        # move to device
        d_A = CudaArray(A)
        d_B = CudaArray(B)
        # compute
        #C = alpha*(A*B.') + conj(alpha)*(B*A.') + beta*C
        C = alpha*(A*B.' + B*A.')
        d_C = CUBLAS.syr2k('U','N',alpha,d_A,d_B)
        # move back to host and compare
        C = triu(C)
        h_C = to_host(d_C)
        h_C = triu(h_C)
        @test C ≈ h_C
    end
end

@testset "her2k!" begin
    @testset for (elty1, elty2) in [(Complex64, Float32), (Complex128, Float64)]
        # generate parameters
        alpha = rand(elty1)
        beta = rand(elty2)
        # generate matrices
        A = rand(elty1,m,k)
        B = rand(elty1,m,k)
        C = rand(elty1,m,m)
        C = C + C'
        # move to device
        d_A = CudaArray(A)
        d_B = CudaArray(B)
        d_C = CudaArray(C)
        # compute
        #C = alpha*(A*B') + conj(alpha)*(B*A') + beta*C
        C = alpha*(A*B') + conj(alpha)*(B*A') + beta*C
        CUBLAS.her2k!('U','N',alpha,d_A,d_B,beta,d_C)
        # move back to host and compare
        C = triu(C)
        h_C = to_host(d_C)
        h_C = triu(h_C)
        @test C ≈ h_C
    end
end

@testset "her2k" begin
    @testset for elty in [Complex64, Complex128]
        # generate matrices
        A = rand(elty,m,k)
        B = rand(elty,m,k)
        # move to device
        d_A = CudaArray(A)
        d_B = CudaArray(B)
        # compute
        C = A*B' + B*A'
        d_C = CUBLAS.her2k('U','N',d_A,d_B)
        # move back to host and compare
        C = triu(C)
        h_C = to_host(d_C)
        h_C = triu(h_C)
        @test C ≈ h_C
    end
end

@testset "trmm!" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # generate parameter
        alpha = rand(elty)
        # generate matrices
        A = rand(elty,m,m)
        A = triu(A)
        B = rand(elty,m,n)
        C = zeros(elty,m,n)
        # move to device
        d_A = CudaArray(A)
        d_B = CudaArray(B)
        d_C = CudaArray(C)
        # compute
        C = alpha*A*B
        CUBLAS.trmm!('L','U','N','N',alpha,d_A,d_B,d_C)
        # move to host and compare
        h_C = to_host(d_C)
        @test C ≈ h_C
    end
end

@testset "trmm" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # generate parameter
        alpha = rand(elty)
        # generate matrices
        A = rand(elty,m,m)
        A = triu(A)
        B = rand(elty,m,n)
        # move to device
        d_A = CudaArray(A)
        d_B = CudaArray(B)
        # compute
        C = alpha*A*B
        d_C = CUBLAS.trmm('L','U','N','N',alpha,d_A,d_B)
        # move to host and compare
        h_C = to_host(d_C)
        @test C ≈ h_C
    end
end

@testset "trsm!" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # generate parameter
        alpha = rand(elty)
        # generate matrices
        A = rand(elty,m,m)
        A = triu(A)
        B = rand(elty,m,n)
        # move to device
        d_A = CudaArray(A)
        d_B = CudaArray(B)
        # compute
        C = alpha*(A\B)
        CUBLAS.trsm!('L','U','N','N',alpha,d_A,d_B)
        # move to host and compare
        h_C = to_host(d_B)
        @test C ≈ h_C
    end
end

@testset "trsm" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # generate parameter
        alpha = rand(elty)
        # generate matrices
        A = rand(elty,m,m)
        A = triu(A)
        B = rand(elty,m,n)
        # move to device
        d_A = CudaArray(A)
        d_B = CudaArray(B)
        # compute
        C = alpha*(A\B)
        d_C = CUBLAS.trsm('L','U','N','N',alpha,d_A,d_B)
        # move to host and compare
        h_C = to_host(d_C)
        @test C ≈ h_C
    end
end

@testset "trsm_batched!" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # generate parameter
        alpha = rand(elty)
        # generate matrices
        A = [rand(elty,m,m) for i in 1:10]
        map!((x) -> triu(x), A, A)
        B = [rand(elty,m,n) for i in 1:10]
        # move to device
        d_A = CudaArray{elty, 2}[]
        d_B = CudaArray{elty, 2}[]
        for i in 1:length(A)
            push!(d_A,CudaArray(A[i]))
            push!(d_B,CudaArray(B[i]))
        end
        # compute
        CUBLAS.trsm_batched!('L','U','N','N',alpha,d_A,d_B)
        # move to host and compare
        for i in 1:length(d_B)
            C = alpha*(A[i]\B[i])
            h_C = to_host(d_B[i])
            #compare
            @test C ≈ h_C
        end
    end
end

@testset "trsm_batched" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # generate parameter
        alpha = rand(elty)
        # generate matrices
        A = [rand(elty,m,m) for i in 1:10]
        map!((x) -> triu(x), A, A)
        B = [rand(elty,m,n) for i in 1:10]
        # move to device
        d_A = CudaArray{elty, 2}[]
        d_B = CudaArray{elty, 2}[]
        for i in 1:length(A)
            push!(d_A,CudaArray(A[i]))
            push!(d_B,CudaArray(B[i]))
        end
        # compute
        d_C = CUBLAS.trsm_batched('L','U','N','N',alpha,d_A,d_B)
        # move to host and compare
        for i in 1:length(d_C)
            C = alpha*(A[i]\B[i])
            h_C = to_host(d_C[i])
            @test C ≈ h_C
        end
    end
end

@testset "hemm!" begin
    @testset for elty in [Complex64, Complex128]
        # generate parameters
        alpha = rand(elty)
        beta  = rand(elty)
        # generate matrices
        A = rand(elty,m,m)
        A = A + ctranspose(A)
        @test ishermitian(A)
        B = rand(elty,m,n)
        C = rand(elty,m,n)
        # move to device
        d_A = CudaArray(A)
        d_B = CudaArray(B)
        d_C = CudaArray(C)
        # compute
        C = alpha*(A*B) + beta*C
        CUBLAS.hemm!('L','L',alpha,d_A,d_B,beta,d_C)
        # move to host and compare
        h_C = to_host(d_C)
        @test C ≈ h_C
    end
end

@testset "hemm" begin
    @testset for elty in [Complex64, Complex128]
        # generate parameter
        alpha = rand(elty)
        # generate matrices
        A = rand(elty,m,m)
        A = A + ctranspose(A)
        @test ishermitian(A)
        B = rand(elty,m,n)
        # move to device
        d_A = CudaArray(A)
        d_B = CudaArray(B)
        # compute
        C = alpha*(A*B)
        d_C = CUBLAS.hemm('L','U',alpha,d_A,d_B)
        # move to host and compare
        h_C = to_host(d_C)
        @test C ≈ h_C
    end
end

@testset "geam!" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # generate parameters
        alpha = rand(elty)
        beta  = rand(elty)
        # generate matrices
        A = rand(elty,m,n)
        B = rand(elty,m,n)
        C = zeros(elty,m,n)
        # move to device
        d_A = CudaArray(A)
        d_B = CudaArray(B)
        d_C = CudaArray(C)
        # compute
        C = alpha*A + beta*B
        CUBLAS.geam!('N','N',alpha,d_A,beta,d_B,d_C)
        # move to host and compare
        h_C = to_host(d_C)
        @test C ≈ h_C

        #test in place versions too
        C = rand(elty,m,n)
        d_C = CudaArray(C)
        C = alpha*C + beta*B
        CUBLAS.geam!('N','N',alpha,d_C,beta,d_B,d_C)
        # move to host and compare
        h_C = to_host(d_C)
        @test C ≈ h_C
        C = rand(elty,m,n)
        d_C = CudaArray(C)
        C = alpha*A + beta*C
        CUBLAS.geam!('N','N',alpha,d_A,beta,d_C,d_C)
        # move to host and compare
        h_C = to_host(d_C)
        @test C ≈ h_C

        #test setting C to zero
        C = rand(elty,m,n)
        d_C = CudaArray(C)
        alpha = zero(elty)
        beta  = zero(elty)
        CUBLAS.geam!('N','N',alpha,d_A,beta,d_B,d_C)
        h_C = to_host(d_C)
        @test h_C ≈ zeros(elty,m,n)

        # bounds checking
        @test_throws DimensionMismatch CUBLAS.geam!('N','T',alpha,d_A,beta,d_B,d_C)
        @test_throws DimensionMismatch CUBLAS.geam!('T','T',alpha,d_A,beta,d_B,d_C)
        @test_throws DimensionMismatch CUBLAS.geam!('T','N',alpha,d_A,beta,d_B,d_C)
    end
end

@testset "geam" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # generate parameter
        alpha = rand(elty)
        beta  = rand(elty)
        # generate matrices
        A = rand(elty,m,n)
        B = rand(elty,m,n)
        # move to device
        d_A = CudaArray(A)
        d_B = CudaArray(B)
        C = zeros(elty,m,n)
        # compute
        C = alpha*A + beta*B
        d_C = CUBLAS.geam('N','N',alpha,d_A,beta,d_B)
        # move to host and compare
        h_C = to_host(d_C)
        @test C ≈ h_C
    end
end

@testset "getrf_batched!" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # generate matrices
        A = [rand(elty,m,m) for i in 1:10]
        # move to device
        d_A = CudaArray{elty, 2}[]
        for i in 1:length(A)
            push!(d_A,CudaArray(A[i]))
        end
        pivot, info = CUBLAS.getrf_batched!(d_A, false)
        h_info = to_host(info)
        for As in 1:length(d_A)
            C   = lufact!(copy(A[As]), Val{false}) # lufact(A[As],pivot=false)
            h_A = to_host(d_A[As])
            #reconstruct L,U
            dL = eye(elty,m)
            dU = zeros(elty,(m,m))
            k = h_info[As]
            if( k >= 0 )
                dL += tril(h_A,-k-1)
                dU += triu(h_A,k)
            end
            #compare
            @test isapprox(C[:L], dL, rtol=1e-2)
            @test isapprox(C[:U], dU, rtol=1e-2)
        end
        for i in 1:length(A)
            d_A[ i ] = CudaArray(A[i])
        end
        pivot, info = CUBLAS.getrf_batched!(d_A, true)
        h_info = to_host(info)
        h_pivot = to_host(pivot)
        for As in 1:length(d_A)
            C   = lufact(A[As])
            h_A = to_host(d_A[As])
            #reconstruct L,U
            dL = eye(elty,m)
            dU = zeros(elty,(m,m))
            k = h_info[As]
            if( k >= 0 )
                dL += tril(h_A,-k-1)
                dU += triu(h_A,k)
            end
            #compare pivots
            @test length(setdiff(h_pivot[:,As],C[:p])) == 0
            #make device pivot matrix
            P = eye(m)
            for row in 1:m
                temp = copy(P[row,:])
                P[row,:] = P[h_pivot[row,As],:]
                P[h_pivot[row,As],:] = temp
            end
            @test inv(P)*dL*dU ≈ inv(C[:P]) * C[:L] * C[:U]
        end
    end
end

@testset "getrf_batched" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # generate matrices
        A = [rand(elty,m,m) for i in 1:10]
        # move to device
        d_A = CudaArray{elty, 2}[]
        for i in 1:length(A)
            push!(d_A,CudaArray(A[i]))
        end
        pivot, info, d_B = CUBLAS.getrf_batched(d_A, false)
        h_info = to_host(info)
        for Bs in 1:length(d_B)
            C   = lufact!(copy(A[Bs]),Val{false}) # lufact(A[Bs],pivot=false)
            h_B = to_host(d_B[Bs])
            #reconstruct L,U
            dL = eye(elty,m)
            dU = zeros(elty,(m,m))
            k = h_info[Bs]
            if( h_info[Bs] >= 0 )
                dU += triu(h_B,k)
                dL += tril(h_B,-k-1)
            end
            #compare
            @test isapprox(C[:L], dL, rtol=1e-2)
            @test isapprox(C[:U], dU, rtol=1e-2)
        end
    end
end

@testset "getri_batched" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # generate matrices
        A = [rand(elty,m,m) for i in 1:10]
        # move to device
        d_A = CudaArray{elty, 2}[]
        for i in 1:length(A)
            push!(d_A,CudaArray(A[i]))
        end
        pivot, info = CUBLAS.getrf_batched!(d_A, true)
        h_info = to_host(info)
        for Cs in 1:length(h_info)
            @test h_info[Cs] == 0
        end
        pivot, info, d_C = CUBLAS.getri_batched(d_A, pivot)
        h_info = to_host(info)
        for Cs in 1:length(d_C)
            C   = inv(A[Cs])
            h_C = to_host(d_C[Cs])
            @test h_info[Cs] == 0
            @test C ≈ h_C
        end
    end
end

@testset "matinv_batched" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # generate matrices
        A = [rand(elty,m,m) for i in 1:10]
        # move to device
        d_A = CudaArray{elty, 2}[]
        for i in 1:length(A)
            push!(d_A,CudaArray(A[i]))
        end
        info, d_C = CUBLAS.matinv_batched(d_A)
        for Cs in 1:length(d_C)
            C   = inv(A[Cs])
            h_C = to_host(d_C[Cs])
            @test C ≈ h_C
        end
    end
end

@testset "geqrf_batched!" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # generate matrices
        A = [rand(elty,m,n) for i in 1:10]
        # move to device
        d_A = CudaArray{elty, 2}[]
        for i in 1:length(A)
            push!(d_A,CudaArray(A[i]))
        end
        tau, d_A = CUBLAS.geqrf_batched!(d_A)
        for As in 1:length(d_A)
            C   = qrfact(A[As])
            h_A = to_host(d_A[As])
            h_tau = to_host(tau[As])
            # build up Q
            Q = eye(elty,min(m,n))
            for i in 1:min(m,n)
                v = zeros(elty,m)
                v[i] = one(elty)
                v[i+1:m] = h_A[i+1:m,i]
                Q *= eye(elty,m) - h_tau[i] * v * v'
            end
            @test Q≈full(C[:Q])
        end
    end
end

@testset "geqrf_batched" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # generate matrices
        A = [rand(elty,m,n) for i in 1:10]
        # move to device
        d_A = CudaArray{elty, 2}[]
        for i in 1:length(A)
            push!(d_A,CudaArray(A[i]))
        end
        tau, d_B = CUBLAS.geqrf_batched!(d_A)
        for Bs in 1:length(d_B)
            C   = qrfact(A[Bs])
            h_B = to_host(d_B[Bs])
            h_tau = to_host(tau[Bs])
            # build up Q
            Q = eye(elty,min(m,n))
            for i in 1:min(m,n)
                v = zeros(elty,m)
                v[i] = one(elty)
                v[i+1:m] = h_B[i+1:m,i]
                Q *= eye(elty,m) - h_tau[i] * v * v'
            end
            @test Q≈full(C[:Q])
        end
    end
end

@testset "gels_batched!" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # generate matrices
        A = [rand(elty,n,n) for i in 1:10]
        C = [rand(elty,n,k) for i in 1:10]
        # move to device
        d_A = CudaArray{elty, 2}[]
        d_C = CudaArray{elty, 2}[]
        for i in 1:length(A)
            push!(d_A,CudaArray(A[i]))
            push!(d_C,CudaArray(C[i]))
        end
        d_A, d_C, info = CUBLAS.gels_batched!('N',d_A, d_C)
        for Cs in 1:length(d_C)
            X = A[Cs]\C[Cs]
            h_C = to_host(d_C[Cs])
            @test X≈h_C
        end
    end
end

@testset "gels_batched" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # generate matrices
        A = [rand(elty,n,n) for i in 1:10]
        C = [rand(elty,n,k) for i in 1:10]
        # move to device
        d_A = CudaArray{elty, 2}[]
        d_C = CudaArray{elty, 2}[]
        for i in 1:length(A)
            push!(d_A,CudaArray(A[i]))
            push!(d_C,CudaArray(C[i]))
        end
        d_B, d_D, info = CUBLAS.gels_batched('N',d_A, d_C)
        for Ds in 1:length(d_D)
            X = A[Ds]\C[Ds]
            h_D = to_host(d_D[Ds])
            @test X ≈ h_D
        end
    end
end

@testset "dgmm!" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # generate matrices
        A = rand(elty,m,n)
        C = rand(elty,m,n)
        X = rand(elty,m)
        # move to device
        d_A = CudaArray(A)
        d_C = CudaArray(C)
        d_X = CudaArray(X)
        # compute
        C = diagm(X) * A
        CUBLAS.dgmm!('L',d_A,d_X,d_C)
        # move to host and compare
        h_C = to_host(d_C)
        @test C ≈ h_C
        # bounds checking
        @test_throws DimensionMismatch CUBLAS.dgmm!('R',d_A,d_X,d_C)
        A = rand(elty,m,m)
        d_A = CudaArray(A)
        @test_throws DimensionMismatch CUBLAS.dgmm!('L',d_A,d_X,d_C)
    end
end

@testset "dgmm" begin
    @testset for elty in [Float32, Float64, Complex64, Complex128]
        # generate matrices
        A = rand(elty,m,n)
        X = rand(elty,m)
        # move to device
        d_A = CudaArray(A)
        d_X = CudaArray(X)
        # compute
        C = diagm(X) * A
        d_C = CUBLAS.dgmm('L',d_A,d_X)
        # move to host and compare
        h_C = to_host(d_C)
        @test C ≈ h_C
    end
end


function normalizedcut(A::AbstractMatrix{T}, b::AbstractVector{T}, k::Int,
                        isfun::Bool, tol::Float64) where {T}
    a = @view A[k, :]
    β = b[k]
    na = norm(a, 2)
    if isfun || na < tol
        a, β
    else
        a / na, β / na
    end
end

"""
    checkredundancy{T}(A::AbstractMatrix{T}, b::AbstractVector{T},
                       Anew::AbstractMatrix{T}, bnew::AbstractVector{T},
                       isfun::Bool, islb::Bool, tol::Float64, ident::Bool=false)

Check redundant cuts between the Polyhedra (A,b) and (Anew, bnew).
Return index of redundant cuts in `Anew`.

# Arguments
* `isfun::Bool`
    States if the Polyhedra defines a function
* `islb::Bool`
    States if the Polyhedra is a lower-bound or an upper-bound
* `tol::Float64`
    Tolerance of redundancy check
* `ident::Bool`
    States whether (A,b)==(Anew,bnew) if we want to remove redundant
    lines in a single Polyhedra

"""
function checkredundancy(A::AbstractMatrix{T}, b::AbstractVector{T},
                            Anew::AbstractMatrix{T}, bnew::AbstractVector{T},
                            isfun::Bool, islb::Bool, tol::Float64,
                            ident::Bool=false) where {T}
    # index of redundants cuts
    redundants = Int[]
    # number of new lines
    nnew = size(Anew, 1)

    for kk in 1:nnew
        a, β = normalizedcut(Anew, bnew, kk, isfun, tol)
        chk, indk = isinside(A, b, a, isfun, tol)
        # if ident is true, we need to take care of the possible match
        # of the cut a with itself in the matrix Anew
        if chk && (~ident || indk!=kk)
            ared, βred = normalizedcut(A, b, indk, isfun, tol)
            if islb ? β <= βred+tol : β+tol >= βred
                push!(redundants, kk)
            end
        end
    end

    redundants
end

"""Check if `λ` is a line of matrix `A`. `λ` might not have the same `eltype`
as `A` and `b` as it might have been scaled by `normalizecut`."""
function isinside(A::AbstractMatrix{T}, b::AbstractVector{T},
                     λ::AbstractVector, isfun::Bool, tol::Float64) where {T}
    nlines = size(A, 1)

    check = false
    k = 0
    while ~check && k < nlines
        k += 1
        a, β = normalizedcut(A, b, k, isfun, tol)
        check = (norm(a - λ, Inf) < tol)
    end
    check, k
end

checkredundancy(A, b, isfun, islb, tol)=checkredundancy(A, b, A, b, isfun, islb, tol, true)


#*******************************************************/
#* Copyright(c) 2018 by Artelys                        */
#* This source code is subject to the terms of the     */
#* MIT Expat License (see LICENSE.md)                  */
#*******************************************************/

#*++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
#*  This example demonstrates how to use Knitro to solve the following
#*  simple problem with a second order cone constraint.
#*
#*  min   x2-1 + x0^2 + x1^2 +(x2+x3)^2
#*  s.t.  sqrt(x0^2 +(2*x2)^2)  - 10*x1 <= 0 (c0)
#*        x3^2 + 5*x0 <= 100                  (c1)
#*        2*x1 + 3*x2 <= 100                  (c2)
#*        x2 <= 1, x1 >= 1, x3 >= 2
#*
#*  Note that the first constraint c0 is a second order cone
#*  constraint that can be written in the form: ||Ax+b||<=c'x
#*  where A = [1, 0, 0, 0 , (b is empty).
#*             0, 0, 2, 0]
#*
#*++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++*/

using KNITRO

function example_conic(; verbose=true)
    #** Create a new Knitro solver instance. */
    kc = KNITRO.KN_new()

    #** Initialize Knitro with the problem definition. */

    ####Add the variables and set their bounds.
    ####*  Note: any unset lower bounds are assumed to be
    ####*  unbounded below and any unset upper bounds are
    ####*  assumed to be unbounded above. */
    n = 4
    KNITRO.KN_add_vars(kc, n)

    xLoBnds = [-KNITRO.KN_INFINITY, 1.0, -KNITRO.KN_INFINITY, 2.0]
    xUpBnds = [KNITRO.KN_INFINITY, KNITRO.KN_INFINITY, 1.0, KNITRO.KN_INFINITY]
    KNITRO.KN_set_var_lobnds(kc, xLoBnds)
    KNITRO.KN_set_var_upbnds(kc, xUpBnds)

    #** Add the constraints and set the RHS and coefficients */
    m = 3
    KNITRO.KN_add_cons(kc, m)
    KNITRO.KN_set_con_upbnd(kc, 0, 0.0)
    KNITRO.KN_set_con_upbnd(kc, 1, 100.0)
    KNITRO.KN_set_con_upbnd(kc, 2, 100.0)

    #** coefficients for linear terms in constraint c2 */
    indexVars1 = Cint[1, 2]
    coefs1 = [2.0, 3.0]
    KNITRO.KN_add_con_linear_struct(kc, 2, indexVars1, coefs1)

    #** coefficient for linear term in constraint c1 */
    indexVars2 = Cint[0]
    coefs2 = [5.0]
    KNITRO.KN_add_con_linear_struct(kc, 1, indexVars2, coefs2)

    #** coefficient for linear term in constraint c0 */
    indexVars3 = Cint[1]
    coefs3 = [-10.]
    KNITRO.KN_add_con_linear_struct(kc, 0, indexVars3, coefs3)

    #** coefficient for quadratic term in constraint c1 */
    qconIndexVar1 = 3;
    qconIndexVar2 = 3;
    qconCoef = 1.0;
    KNITRO.KN_add_con_quadratic_struct(kc, 1, qconIndexVar1, qconIndexVar2, qconCoef)

    #** Coefficients for L2-norm constraint components in c0.
    #*  Assume the form ||Ax+b|| (here with b = 0)
    dimA = 2;   # A = [1, 0, 0, 0; 0, 0, 2, 0] has two rows */
    nnzA = 2;
    indexRowsA = Cint[0, 1]
    indexVarsA = Cint[0, 2]
    coefsA = [1., 2.]
    b = [0., 0.]
    KNITRO.KN_add_con_L2norm(kc, 0, dimA, nnzA, indexRowsA, indexVarsA, coefsA, b)

    #* Set minimize or maximize(if not set, assumed minimize) */
    KNITRO.KN_set_obj_goal(kc, KNITRO.KN_OBJGOAL_MINIMIZE)

    #** Add constant value to the objective. */
    KNITRO.KN_add_obj_constant(kc, -1.0)

    #** Set quadratic objective structure.
    #*  Note:(x2 + x3)^2 = x2^2 + 2*x2*x3 + x3^2 */
    qobjIndexVars1 = Cint[0, 2, 3, 2, 1]
    qobjIndexVars2 = Cint[0, 2, 3, 3, 1]
    qobjCoefs = [1., 1., 1., 2., 1.]
    KNITRO.KN_add_obj_quadratic_struct(kc, qobjIndexVars1, qobjIndexVars2, qobjCoefs)

    #** Add linear objective term. */
    lobjIndexVar = Cint[2]
    lobjCoef = [1.0]
    KNITRO.KN_add_obj_linear_struct(kc, lobjIndexVar, lobjCoef)

    #** Interior/Direct algorithm is required for models with
    #*  L2 norm structure.
    KNITRO.KN_set_param(kc, KNITRO.KN_PARAM_ALGORITHM, KNITRO.KN_ALG_BAR_DIRECT)
    #** Enable the special barrier tools for second order cone(SOC) constraints. */
    KNITRO.KN_set_param(kc, KNITRO.KN_PARAM_BAR_CONIC_ENABLE, KNITRO.KN_BAR_CONIC_ENABLE_SOC)
    #** Specify maximum output */
    outlev = verbose ? KNITRO.KN_OUTLEV_ALL : KNITRO.KN_OUTLEV_NONE
    KNITRO.KN_set_param(kc, KNITRO.KN_PARAM_OUTLEV, outlev)
    #** Specify special barrier update rule */
    KNITRO.KN_set_param(kc, KNITRO.KN_PARAM_BAR_MURULE, KNITRO.KN_BAR_MURULE_FULLMPC)

    #** Solve the problem.
    ####*
    ####*  Return status codes are defined in "knitro.h" and described
    ####*  in the Knitro manual. */
    nStatus = KNITRO.KN_solve(kc)
    nStatus, objSol, x, _ = KNITRO.KN_get_solution(kc)
    feasError = KNITRO.KN_get_abs_feas_error(kc)
    optError = KNITRO.KN_get_abs_opt_error(kc)

    #** An example of obtaining solution information. */
    if verbose
        println("Knitro converged with final status = ", nStatus)
        println("  optimal objective value  = ", objSol)
        println("  optimal primal values x  = ", x)
        println("  feasibility violation    = ", feasError)
        println("  KKT optimality violation = ", optError)
    end

    #** Delete the Knitro solver instance. */
    KNITRO.KN_free(kc)
end

example_conic(; verbose=isdefined(Main, :KN_VERBOSE) ? KN_VERBOSE : true)


include("euler/euler.jl")
using .BigIntegers: BigInteger
using BenchmarkTools
BenchmarkTools.DEFAULT_PARAMETERS.samples = 100

compute()::BigInteger = 28_433 * powermod(BigInteger(2), 7_830_457, 10_000_000_000) % 10_000_000_000 + 1

compute()

@benchmark compute()


# KKT QP with Equality Constraints Solver

include("../auglag/auglag-core.jl")

using SparseArrays

@doc raw"""
    solveQP_AffineEq

We have a problem of the form:

```math
\underset{x}{\text{minimize}} \quad \frac{1}{2} x^⊤Q x + p^⊤x
```
```math
\text{subject to} \quad Ax = b
```

This yields the KKT system:

```math
[Q A^⊤; A 0] \ [x; λ] = [-p; b]
```

This can be solved with LU Factorization, but we will let Julia's "\" function
handle that.

"""
function solveQP_AffineEq(Q, p, A, b)
    hMat = [Q A'; A spzeros(size(A, 1), size(A, 1))]
    dMat = [-p; b]

    return hMat \ Array(dMat)
end

function solveQP_AffineEq(obj::objectiveQP, constraints::AL_AffineEquality)
    Q = obj.Q
    p = obj.p

    A = constraints.A
    b = constraints.b

    return solveQP_AffineEq(Q, p, A, b)
end


function cox_f(S, fs, ls, ξ , X, β, λ, alive, afterΘ) # preprocessed already
    Xβ = X*β
    Θ = exp.(Xβ)
    after!(afterΘ, Θ)

    #compute loglikelihood
    y = 0.
    for i in 1:length(fs)
        for j in (fs[i]):(ls[i])
            ρ = (alive[j]-alive[fs[i]])/(alive[fs[i]]-alive[ls[i]+1])
            ϕ = afterΘ[fs[i]]-ρ*(afterΘ[fs[i]]-afterΘ[ls[i]+1])
            y -= Xβ[j] -log(ϕ)
        end
    end
    y += λ*(β' * β)
    return y
end

# preprocessed already:
# fs = index first deaths, ls = index last deaths,
# X is covariates, ξ is covariate covariate transpose
function cox_h!(grad,hes, S, fs, ls, ξ , X, β, λ, alive, afterΘ, afterXΘ, afterξΘ)
    #compute relevant quantities for loglikelihood, score, fischer_info

    Xβ = X*β
    Θ = exp.(Xβ)
    after!(afterΘ,Θ)
    after!(afterXΘ, X.*Θ)
    after!(afterξΘ,ξ.*Θ)

    #compute loglikelihood, score, fischer_info
    #From v0.6 remember to use . notation and do it inplace, possibly using views!
    y = 0.
    grad[:] = 0.
    hes[:] = 0.

    # preallocate
    Z = zeros(size(X,2))
    Ξ = zeros(size(X,2),size(X,2))

    for i in 1:length(fs)
        for j in (fs[i]):(ls[i])
            ρ = (alive[j]-alive[fs[i]])/(alive[fs[i]]-alive[ls[i]+1])
            ϕ = afterΘ[fs[i]]-ρ*(afterΘ[fs[i]]-afterΘ[ls[i]+1])
            for k in eachindex(Z)
                @inbounds Z[k] = afterXΘ[fs[i],k]-ρ*(afterXΘ[fs[i],k]-afterXΘ[ls[i]+1,k])
            end
            for k2 in 1:size(Ξ,2)
                for k1 in 1:size(Ξ,1)
                    @inbounds Ξ[k1,k2] = afterξΘ[fs[i],k1,k2]-ρ*(afterξΘ[fs[i],k1,k2]-afterξΘ[ls[i]+1,k1,k2])
                end
            end
            y -= Xβ[j] -log(ϕ)
            for k2 in 1:size(Ξ,2)
                @inbounds grad[k2] -= X[j,k2]
                @inbounds grad[k2] += Z[k2]/ϕ
                for k1 in 1:size(Ξ,1)
                    @inbounds hes[k1, k2] += Ξ[k1,k2]/ϕ - Z[k1]*Z[k2]/ϕ^2
                end
            end
        end
    end
    y += λ*(β' * β)

    for k1 in 1:size(Ξ,1)
        grad[k1] +=  2*λ*β[k1]
        hes[k1,k1] +=  2*λ
    end
    return y
end

function coxph(S::AbstractVector,X::AbstractArray; l2_cost = 0., kwargs...)
    ξ = zeros(size(X,1),size(X,2),size(X,2))
    for k2 in 1:size(ξ,3)
        for k1 in 1:size(ξ,2)
            for i in 1:size(ξ,1)
                @inbounds ξ[i,k1,k2] = X[i,k1]*X[i,k2]
            end
        end
    end

    # compute first and last!

    fs = find(firsts(S))
    ls = find(lasts(S))

    # do optimization

    alive = after(ones(Int64, length(S)))
    afterΘ = init_after(zeros(size(X,1)))
    afterXΘ = init_after(X.*zeros(size(X,1)))
    afterξΘ = init_after(ξ.*zeros(size(X,1)))

    f1 = (β) -> cox_f(S, fs, ls, ξ , X, β, l2_cost, alive, afterΘ)
    h1! = (β,grad,hes) -> cox_h!(grad,hes, S, fs, ls, ξ , X, β, l2_cost,alive, afterΘ, afterXΘ, afterξΘ)
    return newton_raphson(f1,h1!, zeros(size(X,2)); kwargs...)
end

function coxph(formula::Formula, data::DataFrame; l2_cost = 0., kwargs...)
    sorted_data = sort(data, cols = formula.lhs)
    M = DataFrames.ModelFrame(formula,sorted_data)
    S = convert(Array, M.df[:,1])
    model_matrix = DataFrames.ModelMatrix(M)
    X = model_matrix.m[:,2:size(model_matrix.m,2)]
    β, neg_ll,grad, hes =  coxph(S, X; l2_cost = l2_cost, kwargs...)
    rownms = coefnames(M)[2:end]
    se = sqrt.(diag(pinv(hes)))
    z_score = β./se
    pvalues = 2*cdf(Normal(),-abs.(z_score))
    coefmat = CoefTable(hcat([β, se, z_score, pvalues]...),
    ["Estimate", "Std.Error", "z value", "Pr(>|z|)"], rownms, 4)
    CoxModel("Cox; ", formula, coefmat, M, -neg_ll, -grad, hes)
end


#=

  Set covering problem in Julia ConstraintSolver.jl

  Problem from
  Katta G. Murty: "Optimization Models for Decision Making", page 302f
  http://ioe.engin.umich.edu/people/fac/books/murty/opti_model/junior-7.pdf

  10 senators making a committee, where there must at least be one
  representative from each group:
  group:        senators:
  southern      1 2 3 4 5
  northern      6 7 8 9 10
  liberals      2 3 8 9 10
  conservative  1 5 6 7
  democrats     3 4 5 6 7 9
  republicans   1 2 8 10

  The objective is to minimize the number of senators.

  Model created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#
using ConstraintSolver, JuMP
using Cbc, GLPK, Ipopt
const CS = ConstraintSolver
include("constraints_utils.jl")

function set_covering3(problem,print_solutions=true,all_solutions=true)

    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    glpk_optimizer = optimizer_with_attributes(GLPK.Optimizer)
    ipopt_optimizer = optimizer_with_attributes(Ipopt.Optimizer)

    model = Model(optimizer_with_attributes(CS.Optimizer,   # "all_solutions"=> all_solutions,
                                                            "all_optimal_solutions"=>all_solutions,
                                                            "logging"=>[],

                                                            "traverse_strategy"=>:BFS,
                                                            # "traverse_strategy"=>:DFS,
                                                            # "traverse_strategy"=>:DBFS,

                                                            # "branch_split"=>:Smallest,
                                                            # "branch_split"=>:Biggest,
                                                            "branch_split"=>:InHalf,

                                                            # https://wikunia.github.io/ConstraintSolver.jl/stable/options/#branch_strategy-(:Auto)
                                                            "branch_strategy" => :IMPS, # default
                                                            # "branch_strategy" => :ABS, # Activity Based Search
                                                            # "activity.decay" => 0.999, # default 0.999
                                                            # "activity.max_probes" => 100, # default, 10
                                                            # "activity.max_confidence_deviation" => 20, # default 20

                                                            # "simplify"=>false,
                                                            # "simplify"=>true, # default

                                                            "time_limit"=>6,

                                                            # "backtrack" => false, # default true
                                                            # "backtrack_sorting" => false, # default true

                                                            # "lp_optimizer" => cbc_optimizer,
                                                            "lp_optimizer" => glpk_optimizer,
                                                            # "lp_optimizer" => ipopt_optimizer,
                                        ))
    belongs = problem[:belongs]
    num_groups,num_senators = size(belongs)

    # Which senator to select
    @variable(model, x[1:num_senators], Bin)
    @variable(model, 0 <= min_val <= num_senators, Int)

    # cover all groups with senators
    for i in 1:num_groups
        @constraint(model, sum(x.*belongs[i,:])>=1)
    end

    @constraint(model, min_val == sum(x))

    @objective(model,Min,min_val)

    # Solve the problem
    optimize!(model)

    status = JuMP.termination_status(model)
    # println("status:$status")
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        println("num_sols:$num_sols\n")
        if print_solutions
            for sol in 1:num_sols
                println("solution #$sol")
                x_val = convert.(Integer,JuMP.value.(x; result=sol))
                min_val_val = convert.(Integer,JuMP.value.(min_val; result=sol))
                println("min_val:$min_val_val")
                println("x:$x_val")
                println("senators selected:$([i for i in 1:num_senators if x_val[i] == 1])")
                println()
            end
        end
    else
        println("status:$status")
    end

    return status
end

problem = Dict(
    :belongs => resize_matrix([
        [1, 1, 1, 1, 1, 0, 0, 0, 0, 0],   # 1 southern
        [0, 0, 0, 0, 0, 1, 1, 1, 1, 1],   # 2 northern
        [0, 1, 1, 0, 0, 0, 0, 1, 1, 1],   # 3 liberals
        [1, 0, 0, 0, 1, 1, 1, 0, 0, 0],   # 4 conservative
        [0, 0, 1, 1, 1, 1, 1, 0, 1, 0],   # 5 democrats
        [1, 1, 0, 0, 0, 0, 0, 1, 0, 1]    # 6 republicans
    ]                                   )
)

@time set_covering3(problem,true,true)


import Pkg; Pkg.activate("../")
using Turing: Bijectors

b = Bijectors.SimplexBijector()
ib = inv(b)
x = [.3, .5, .2]
y = b(x)

Bijectors.logabsdetjac(b, x)
Bijectors.logabsdetjac(ib, y)


"""
    SpectralNeuralOperator{F,L,F1,F2,P}

A type representing a Neural Operator whose forward pass is of the form

    y(t) = σ((B*x)(t) - λ(t)*(B*x)(t) + b(t))

where `x`, `λ` and `b` are functions, and `B` is a linear operator.

By default, `y` is projected in the Chebyshev polynomial basis before outputing. This is done to avoid computing multiple nested integrals when stacking multiple layers.

## Fields
 * `B::L` A linear operator on functions ``[0,1] → R``, generally a kernel operator.
 * `λ::F1` a function ``[0,1] → R``
 * `b::F2` a function ``[0,1] → R`` used as bias
 * `σ::F` the activation function
 * `project::P` a projection operator used on the output of the layer. Since by default, the output is a lazy function, this is meant to
"""
struct SpectralNeuralOperator{F,L,F1,F2,P}
    B::L # Linear Operator
    λ::F1 # Spectral Scaling Function
    b::F2 # Bias function
    σ::F # activation function
    project::P # Projection operator
end

SpectralNeuralOperator(B, λ, b) = SpectralNeuralOperator(B, λ, b, identity, f -> Fun(f, 0..1))

function Base.show(io::IO, ::MIME"text/plain", x::SpectralNeuralOperator{F,L,F1,F2,P}) where {F,L,F1,F2,P}
    println(io, "SpectralNeuralOperator with\n kernel type $L\n scaling type $F1\n bias type $F2\n activation function $F\n projection operator $P")
end

Flux.@functor SpectralNeuralOperator

function kernel_fun(sno::SpectralNeuralOperator{F,<:KernelOperator,F1,F2,P}) where {F1,F2,F,P}
    return function k_fun(t,s)
        (1 - sno.λ(t))*sno.B.kernel(t,s)
    end
end

function (sno::SpectralNeuralOperator)(x)
    B, λ, b, σ = sno.B, sno.λ, sno.b, sno.σ
    y = B(x)
    z = t -> σ((1- λ(t))*y(t) + b(t))
    sno.project(z)
end

function b_pullback(sno::SpectralNeuralOperator{F,L,F1,F2,P}, δz, y) where {F,L,F1,F2<:ParametricFun,P}
    dσ(u) = ForwardDiff.derivative(sno.σ, u)
    function fun(u)
        dσ(y(u)) .* to_grad(sno.b, u)
    end
    δb = product_integral(fun, δz)
    return from_grad(sno.b, δb)
end

# default behaviour, no tangent
b_pullback(sno, δz, δy) = ZeroTangent()

function λ_pullback(sno::SpectralNeuralOperator{F,L,F1,F2,P}, δz, y) where {F,L,F1<:ParametricFun,F2,P}
    dσ(u) = ForwardDiff.derivative(sno.σ, u)
    function fun(u)
        dσ(y(u)) .* to_grad(sno.λ, u)
    end
    δλ =  kernel_double_integral(kernel_fun(sno), fun, δz)
    return from_grad(sno.λ, δλ)
end

λ_pullback(sno, δz, δy) = ZeroTangent()

function x_pullback(sno::SpectralNeuralOperator, δz, y)
    dσ(u) = ForwardDiff.derivative(sno.σ, u)
    function δx(s)
        dσ(y(s)) * kernel_transpose_integral(kernel_fun(sno), δz, s)
    end
    return sno.project(δx)
end

function ChainRulesCore.rrule(sno::SpectralNeuralOperator, x)
    Bx = sno.B(x)
    y = sno.project(t -> (1 - sno.λ(t))*Bx(t) + sno.b(t))
    z = t -> sno.σ( (1 - sno.λ(t)) * y(t) + sno.b(t) )
    function sno_pullback(δz)
        δB = ZeroTangent()
        δλ = λ_pullback(sno, δz, y)
        δb = b_pullback(sno, δz, y)
        δσ = NoTangent()
        δproject = NoTangent()
        δx = x_pullback(sno, δz, y)
        return Tangent{SpectralNeuralOperator}(;B=δB, λ=δλ, b=δb, σ=δσ, project=δproject), δx
    end
    return sno.project(z), sno_pullback
end


function PhaseShift(x,angle)

	nx = length(x)
	nf = 2*nx
	nw = int(nf/2) + 1
	X = fft([x;zeros(typeof(x[1]),nf-nx)])
	for iw=1:nw
		X[iw] *= exp(-1im*angle)
	end
	# symmetries
	for iw=nw+1:nf
		X[iw] = conj(X[nf-iw+2])
	end
	x = real(ifft(X,1))
	return x[1:nx]

end


@testset "Homology" begin
    cplx = SimplicialComplex(Simplex(1,2,3),
                        Simplex(2,4),
                        Simplex(3,4),
                        Simplex(5,4),
                        Simplex(6))

    h = homology(Int, cplx)
    @test eltype(h) == Tuple{Int64,Int64,Int64}
    @test grouptype(supertype(typeof(h))) == Nothing
    @test grouptype(typeof(h)) == Int
    @test length(h) == 2

    b, t, _ = group(h,0);
    @test b == 2
    @test t == 0

    itm, st = iterate(h)
    @test itm[1] == 0
    @test itm[2] == 2
    @test st[1] == 1

    itm, st = iterate(h, st)
    @test itm[1] == 1
    @test itm[2] == 1
    @test st[1] == 2

    @test iterate(h, st) === nothing

    @test betti(h) == (2,1)
    @test euler(betti(h)) == 1

    g = withgenerators(h)
    @test eltype(g) == Tuple{Int64,Int64,Int64,Dict{Chain,Int64}}
    @test length(g) == 2

    itm, st = iterate(g)
    @test itm[1] == 0
    @test itm[2] == 2
    @testset for (ch, d) in itm[4]
        @test d == 0
        @test all(k->k == 1, values(ch))
        @test all(k->k ∈ map(hash, map(Simplex, [5,6])), keys(ch))
    end

    itm, st = iterate(g, st)
    @test itm[1] == 1
    @test itm[2] == 1
    @test first(itm[4])[2] == 0
    gn =  Chain(1, Int) + (hash(Simplex(3,4)),1) + (hash(Simplex(2,4)),-1) + (hash(Simplex(2,3)),1)
    @testset for (a,b) in zip( simplify(first(itm[4])[1]), gn)
        @test a == b
    end

    @test iterate(g, st) === nothing

    @test length(generators(g)) == 2
end


module TestHill
using Test
using LinearAlgebra

# Checking Hill numbers
using Diversity
using Diversity.Hill

numspecies = 100;
numcommunities = 8;
manyweights = rand(numspecies, numcommunities);
manyweights *= Diagonal(reshape(mapslices(v -> 1. / sum(v), manyweights, dims=1),
                             (numcommunities)));

@testset "Hill numbers" begin
    for i in 1:size(manyweights, 2)
        @test hillnumber(manyweights[:,i], [0])[!,:diversity] ≈ [numspecies]
    end
end

end


using Revise
using Soss
using Random

# Build the model

m = @model Prior,x,σ begin
    β ~ Prior
    yhat = β .* x
    y ~ For(eachindex(x)) do j
        Normal(yhat[j], σ)
    end
end

# Sample from generative model

Random.seed!(42);

x = randn(3);
Prior = Normal()
args = (Prior=Prior, x=x, σ=2.0);
truth = rand(m(args));

pairs(truth)

# Joint Distributions

pairs(args)

m(args)

# Evaluate log density

pairs(truth)

logpdf(m(args), truth)

# Sample from posterior

pairs(args)

post = dynamicHMC(m(args), (y=truth.y,));
particles(post)

# Determine predictive distribution

pred = predictive(m, :β)

# Sample from posterior predictive distribution

particles(post)

argspost = merge(args, particles(post));
pairs(argspost)

postpred = pred(argspost) |> rand;
pairs(postpred)

q = @model λ begin
β ~ Normal(λ,1)
end

# imp(m(args),q(λ=1.0),truth)
# importanceSample(m(args),q(λ=1.0),truth)

# particles(m(args))


function find_tsp_tour(Ct::Matrix{Float64})
    scale_factor = 1000
    dist_mtx = round.(Int, Ct .* scale_factor)

    tsp_tour, tsp_tour_len = Concorde.solve_tsp(dist_mtx)

    @assert tsp_tour[1] == 1

    return tsp_tour
end

function find_tsp_tour(x::Vector{Float64}, y::Vector{Float64})
    scale_factor = 1000

    tsp_tour, tsp_tour_len = Concorde.solve_tsp(x * scale_factor, y * scale_factor; dist="EUC_2D")

    @assert tsp_tour[1] == 1

    return tsp_tour
end

function exact_partitioning(initial_tour, Ct, Cd; flying_range=MAX_DRONE_RANGE)
    n, _ = size(Ct)

    r = initial_tour
    T = fill(Inf, n, n)
    M = fill(-99, n, n)
    for i in 1:n-1
        for j in i+1:n
            if j == i + 1
                T[i, j] = Ct[r[i], r[j]]
                M[r[i], r[j]] = -1
            else
                for k in i+1:j-1
                    Tk1 = Cd[r[i], r[k]] + Cd[r[k], r[j]]
                    if Tk1 <= flying_range
                        Tk2 = sum([Ct[r[l], r[l+1]] for l in i:k-2]) +
                                Ct[r[k-1], r[k+1]] +
                                sum([Ct[r[l], r[l+1]] for l in k+1:j-1])
                        Tk = max(Tk1, Tk2)
                        if Tk < T[i, j]
                            T[i, j] = Tk
                            M[r[i], r[j]] = r[k]
                        end
                    end
                end

            end
        end
    end

    V = zeros(n)
    P = fill(-1, n)

    V[1] = 0
    for i in 2:n
        VV = [V[k] + T[k, i] for k in 1:i-1]
        V[i] = minimum(VV)
        P[i] = r[argmin(VV)]
    end

    # Retrieving solutions.
    combined_nodes = Int[]
    current_idx = n
    current = r[current_idx]
    while current != -1
        pushfirst!(combined_nodes, current)
        current_idx = findfirst(x -> x == current, r)
        current = P[current_idx]
    end
    drone_only_nodes = Int[]
    drone_route = Int[]
    push!(drone_route, combined_nodes[1])
    @assert combined_nodes[1] == r[1]
    for i in 1:length(combined_nodes)-1
        j1 = combined_nodes[i]
        j2 = combined_nodes[i+1]
        if M[j1, j2] != -1
            push!(drone_only_nodes, M[j1, j2])
            push!(drone_route, M[j1, j2])
        end
        push!(drone_route, j2)
    end
    truck_route = setdiff(initial_tour, drone_only_nodes)

    obj_val = objective_value(truck_route, drone_route, Ct, Cd)
    final_time = V[end]

    @assert isapprox(obj_val, final_time)
    return final_time, truck_route, drone_route
end

function two_point_move(tour, i, j)
    if i >= j
        return tour, false
    end

    tmp = copy(tour)
    tmp[j] = tour[i]
    tmp[i] = tour[j]
    return tmp, true
end

function one_point_move(tour, i, j)
    tmp = copy(tour)
    deleteat!(tmp, i)
    insert!(tmp, j, tour[i])
    return tmp, true
end

function two_opt_move(tour, i, j)
    if i >= j
        return tour, false
    end

    tmp = copy(tour)
    tmp[i:j] = tour[j:-1:i]
    return tmp, true
end

# Main function to call
function tsp_ep_all(
    x_coordinates,
    y_coordinates,
    truck_cost_factor,
    drone_cost_factor;
    local_search_methods=[two_point_move, one_point_move, two_opt_move],
    flying_range=MAX_DRONE_RANGE,
    time_limit=MAX_TIME_LIMIT
)
    """
    Runs `TSP-ep-all` heuristic algorithm of Agatz et al.

    `x_coordinates[1]`, `y_coordinates[1]`: the coordinates of the depot, then followed by all customer location coordinates
    `truck_cost_factor`: as defined in Agatz et al. instances
    `drone_cost_factor`: as defined in Agatz et al. instances
    """

    Ct, Cd = cost_matrices_with_dummy(x_coordinates, y_coordinates, truck_cost_factor, drone_cost_factor)
    n_nodes = length(x_coordinates)
    n1, n2 = size(Ct)
    @assert n_nodes + 1 == n1
    @assert size(Ct) == size(Cd)

    tsp_tour = find_tsp_tour(x_coordinates, y_coordinates)
    push!(tsp_tour, n_nodes+1) # adding a dummy node for the returning depot

    return tsp_ep_all(Ct, Cd, tsp_tour, flying_range=flying_range, local_search_methods=local_search_methods, time_limit=time_limit)
end

function tsp_ep_all(
    Ct,
    Cd,
    init_tour;
    local_search_methods=[two_point_move, one_point_move, two_opt_move],
    flying_range=MAX_DRONE_RANGE,
    time_limit=MAX_TIME_LIMIT
)
    n, _ = size(Ct)

    improved = true
    best_obj, t_route, d_route = exact_partitioning(init_tour, Ct, Cd, flying_range=flying_range)

    best_tour = copy(init_tour)
    best_t_route = copy(t_route)
    best_d_route = copy(d_route)

    if isempty(local_search_methods)
        return best_obj, best_t_route, best_d_route
    end

    time0 = time()

    is_time_over = false

    while improved && !is_time_over
        improved = false
        cur_best_obj = best_obj
        cur_best_tour = copy(best_tour)
        cur_best_t_route = copy(best_t_route)
        cur_best_d_route = copy(best_d_route)

        for i in 2:n-1
            if is_time_over
                break
            end

            for local_search in local_search_methods
                for j in 2:n-1
                    if is_time_over
                        break
                    end

                    new_tour, is_valid = local_search(best_tour, i, j)
                    is_time_over = time() - time0 > time_limit

                    if is_valid
                        ep_time, t_route, d_route = exact_partitioning(new_tour, Ct, Cd, flying_range=flying_range)
                        if ep_time < cur_best_obj
                            cur_best_tour = copy(new_tour)
                            cur_best_t_route = copy(t_route)
                            cur_best_d_route = copy(d_route)
                            cur_best_obj = ep_time
                            improved = true
                        end
                    end
                end
            end
        end

        if improved
            best_obj = cur_best_obj
            best_tour = cur_best_tour
            best_t_route = cur_best_t_route
            best_d_route = cur_best_d_route
        end
    end

    return best_obj, best_t_route, best_d_route
end


module TestStatistics
    using NullableStats
    using Base.Test

    srand(1)

    for N in (10, 100)
        A = rand(N)
        M = rand(Bool, N)
        mu_A = mean(A)
        nmu_A = Nullable(mu_A)
        i = rand(1:N)
        M[i] = true
        j = rand(1:N)
        while j == i
            j = rand(1:N)
        end
        M[j] = false
        X = NullableArray(A)
        Y = NullableArray(A, M)
        J = find(x -> !x, M)
        B = A[J]
        mu_B = mean(B)
        nmu_B = Nullable(mu_B)

        C = rand(N)
        V = WeightVec(C)
        R = rand(Bool, N)
        R[rand(1:N)] = true
        R[j] = false
        # W = WeightVec(NullableArray(C, R))

        K = find(x -> !x, R)
        L = find(i -> (!M[i] & !R[i]), [1:N...])

        # For testing Base.varzm
        D1 = rand(round(Int, N / 2))
        D2 = -1 .* D1
        D = [D1; D2]
        while mean(D) != 0
            D1 = rand(round(Int, N / 2))
            D2 = -1 .* D1
            D = [D1; D2]
        end
        Q = NullableArray(D)
        S = rand(Bool, round(Int, N / 2))
        U = NullableArray(D, [S; S])
        E = [D[find(x->!x, S)]; D[find(x->!x, S)]]

        @test_throws NullException varm(Y, Nullable{Float64}())

        # Test mean
        for skip in (true, false)
            v = mean(X; skipnull=skip)
            @test_approx_eq v.value mean(A)
            @test !v.isnull

            v = mean(Y; skipnull=skip)
            if skip == false
                @test isequal(v, Nullable{Float64}())
            else
                @test_approx_eq v.value mean(B)
                @test !v.isnull
            end

            v = mean(X, V; skipnull=skip)
            @test_approx_eq v.value mean(A, V)
            @test !v.isnull

            v = mean(Y, V; skipnull=skip)
            if skip == false
                @test isequal(v, Nullable{Float64}())
            else
                @test_approx_eq v.value mean(B, WeightVec(C[J]))
                @test !v.isnull
            end

            # Following tests need to wait until WeightVec constructor is
            # implemented for NullableArray argument
            #
            # v = mean(X, W, skipnull=skip)
            # if skip == false
            #     @test isequal(v, Nullable{Float64}())
            # else
            #     @test_approx_eq v.value mean(A[K], WeightVec(C[K]))
            #     @test !v.isnull
            # end
            # v = mean(Y, W, skipnull=skip)
            # if skip == false
            #     @test isequal(v, Nullable{Float64}())
            # else
            #     @test_approx_eq v.value mean(A[L], WeightVec(C[L]))
            #     @test !v.isnull
            # end
        end

        for corr in (true, false), skip in (true, false)
            # Test Base.varzm
            v = Base.varzm(Q, corrected=corr, skipnull=skip)
            @test_approx_eq v.value Base.varzm(D, corrected=corr)
            @test !v.isnull

            v = Base.varzm(U, corrected=corr, skipnull=skip)
            if skip == false
                @test isequal(v, Nullable{Float64}())
            else
                @test_approx_eq v.value Base.varzm(E, corrected=corr)
                @test !v.isnull
            end

            # Test varm, stdm
            for method in (varm, stdm)
                for mu in (mu_A, nmu_A)
                    v = method(X, mu, corrected=corr, skipnull=skip)
                    @test_approx_eq v.value method(A, mu_A, corrected=corr)
                    @test !v.isnull
                end

                for mu in (mu_B, nmu_B)
                    v = method(Y, mu, corrected=corr, skipnull=skip)
                    if skip == false
                        @test isequal(v, Nullable{Float64}())
                    else
                        @test_approx_eq v.value method(B, mu_B, corrected=corr)
                        @test !v.isnull
                    end
                end
            end

            # Test var, std
            for method in (var, std)
                for mu in (nothing, mu_A, nmu_A, :a)
                    if mu == :a
                        @test_throws ErrorException method(X, mean=mu,
                                                               corrected=corr,
                                                               skipnull=skip)
                    else
                        v = method(X, mean=mu, corrected=corr, skipnull=skip)
                        @test_approx_eq v.value method(A, mean=mu_A, corrected=corr)
                        @test !v.isnull
                    end
                end

                for mu in (0, Nullable(0))
                    v = method(X, mean=mu, corrected=corr, skipnull=skip)
                    @test_approx_eq v.value method(A, mean=0, corrected=corr)
                    @test !v.isnull
                end

                for mu in (nothing, mu_B, nmu_B)
                    v = method(Y, mean=mu, corrected=corr, skipnull=skip)
                    if skip == false
                        @test isequal(v, Nullable{Float64}())
                    else
                        @test_approx_eq v.value method(B, mean=mu_B, corrected=corr)
                        @test !v.isnull
                    end
                end

                for mu in (0, Nullable(0))
                    v = method(Y, mean=mu, corrected=corr, skipnull=skip)
                    if skip == false
                        @test isequal(v, Nullable{Float64}())
                    else
                        @test_approx_eq v.value method(B, mean=0, corrected=corr)
                        @test !v.isnull
                    end
                end
            end
        end
    end
end


import Dates
println("Last updated: ", Dates.now(), " (PT)")

# Load environment
import Pkg; Pkg.activate("../../../")

# Import Libraries
using Turing
import Turing.RandomMeasures.DirichletProcess
import Turing.RandomMeasures.ChineseRestaurantProcess
using Distributions
using PyPlot
import Random
using BenchmarkTools
import StatsBase.countmap

# rand(ChineseRestaurantProcess(DirichletProcess(1.0), [155, 10]))

function extract(chain, sym; burn=0)
    tail  = chain[sym].value.data[(burn + 1):end, :, :]
    return dropdims(tail, dims=3)
end

function counts_plot(x; density=false, color="C0", lw=4)
    cm_x = countmap(x)
    number = collect(keys(cm_x))
    count = let
        c = collect(values(cm_x))
        density ? c / sum(c) : c
    end
    plt.vlines(number, 0, count, color=color, lw=lw)
end

# FIXME: Not working???

# Define model
@model infiniteGMM(x) = begin
    nobs = length(x)

    # Hyper-parameters, i.e. concentration parameter and parameters of H.
    alpha = 1.0
    mu0 = 0.0
    sig0 = 1.0

    # Define random measure, e.g. Dirichlet process.
    rpm = DirichletProcess(alpha)

    # Define the base distribution, i.e. expected value of the Dirichlet process.
    H = Normal(mu0, sig0)

    # Latent assignment.
    z = tzeros(Int, nobs)

    # Locations of the infinitely many clusters.
    mu = tzeros(Float64, 0)

    for i in 1:nobs
        # Number of clusters.
        K = maximum(z)
        nk = Vector{Int}(map(k -> sum(z .== k), 1:K))

        # Draw the latent assignment.
        z[i] ~ ChineseRestaurantProcess(rpm, nk)

        # Create a new cluster?
        if z[i] > K
            push!(mu, 0.0)

            # Draw location of new cluster.
            mu[z[i]] ~ H
        end

        # Draw observation.
        x[i] ~ Normal(mu[z[i]], 1.0)
    end
end

# Generate data
Random.seed!(1)
data = vcat(randn(10), randn(10) .- 5, randn(10) .+ 10)
data .-= mean(data)
data /= std(data);

# plt.scatter(1:length(data), sort(data))
plt.hist(data, bins=15);

# Fit model
Random.seed!(2)
iterations = 1000
model_fun = infiniteGMM(data)
chain = sample(model_fun, SMC(), iterations);

samples = chain.value.data[:, 4:end, 1]
z = samples[:, end-30+1:end]
mu = samples[:, 1:end-30];

mu

nclus = [length(unique(z[i, :])) for i in 1:size(z, 1)]
plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
plt.plot(nclus);
plt.subplot(1, 2, 2);
counts_plot(nclus, density=true);
plt.xlabel("Number of clusters")
plt.ylabel("Posterior Probability");

# counts_plot(z[end-2, :])

counts_plot(z[5,:]);

function rand_mu(mu, z, i)
    mu_vec =  filter(x -> !ismissing(x), mu[i, :])
    z = Int(rand(z[i, :]))
    mu_vec[z]
end
post_pred = [rand(Normal(rand_mu(mu, z, i), 1)) for i in 1:size(z, 1)]
plt.hist(post_pred);

mu


using Test
using CLIMAParameters: AbstractEarthParameterSet
using CLIMAParameters.Planet
using CLIMAParameters.Atmos.SubgridScale

@testset "Atmos SubgridScale" begin
    struct EarthParameterSet <: AbstractEarthParameterSet end

    ps = EarthParameterSet()
    # Atmos.SubgridScale
    @test c_1_KASM(ps) ≈ c_a_KASM(ps) * 0.76^2
    @test c_2_KASM(ps) ≈ c_e2_KASM(ps) + 2 * c_1_KASM(ps)
    @test c_3_KASM(ps) ≈ c_a_KASM(ps)^(3 / 2)

end


@testset "SimpleGrid" begin
    @polyvar x y z w
    f = x^2+y^2+1
    F = AmoebaFiber2D(x^2+y^2+1)
    grid = Grid2D((-2, 2), 50)
    generator = PolynomialAmoebas.startvalue_generator(F)
    B = PolynomialAmoebas.simple_grid(F, grid, generator)
    @test any(B.data)
    B = PolynomialAmoebas.amoeba(f, grid=grid, alg=Simple())
    @test any(B.data)

    I = PolynomialAmoebas.imaginary_projection(f, grid=grid, alg=Simple())
    @test any(I.data)

    f = x^2+y^2+z^2+1
    F = AmoebaFiber3D(f)
    grid = Grid3D(xlims=(-2, 2), ylims=(-2, 2), zlims=(-2, 2), res=(21, 21, 21))
    generator = PolynomialAmoebas.startvalue_generator(F)
    B = PolynomialAmoebas.simple_grid(F, grid, generator)
    @test any(B.data)
    B = PolynomialAmoebas.amoeba(f, grid=grid, alg=Simple())
    @test any(B.data)

    f = 1 + z * cis(0.8π) + z^3*w*cis(1.6π) + z*w^2*cis(0.4π) + w*cis(1.2π)
    C = CoamoebaFiber2D(f)
    grid = Grid2D((0.0, 2π), 50)
    generator = PolynomialAmoebas.DomainStartValueGenerator2D((-5, 5), (-5, 5))
    B = simple_grid(C, grid, generator)
    @test any(B.data)
    B = coamoeba(f, alg=Simple())
    @test any(B.data)
end


import ForwardDiff
import Optim

@inline function geometric_median(a::AbstractVector{V};
                                  x0 = zeros(T, length(first(a)))) where V <: AbstractVector{T} where T
    function _sum_of_distances(x)
        return sum(l2_distance.(a, Ref(x)))
    end
    result = Optim.optimize(_sum_of_distances,
                            x0,
                            Optim.LBFGS(); autodiff = :forward)
    minimizer = Optim.minimizer(result)
    return minimizer
end


# Filter response functions

#
# Frequency response of a digital filter
#

using ..DSP: xcorr

"""
    H, w = freqresp(filter)

Frequency response `H` of a `filter` at (normalized) frequencies `w` in
radians/sample for a digital filter or radians/second for an analog filter
chosen as a reasonable default.
"""
function freqresp(filter::FilterCoefficients)
    w = _freqrange(filter)
    return (freqresp(filter, w), w)
end

"""
    freqresp(filter::FilterCoefficients{:z}, w)

Frequency response of digital `filter` at normalized frequency or
frequencies `w` in radians/sample.
"""
freqresp(filter::FilterCoefficients{:z}, w) = _freq(filter).(exp.(im .* w))

"""
    freqresp(filter::FilterCoefficients{:s}, w)

Frequency response of analog `filter` at frequency or frequencies `w` in
radians/second.
"""
freqresp(filter::FilterCoefficients{:s}, w) = _freq(filter).(im .* w)

_freq(filter::FilterCoefficients) = x::Number -> _freq(filter, x)

_freq(filter::FilterCoefficients, x::Number) =
    _freq(convert(PolynomialRatio, filter), x)
_freq(filter::PolynomialRatio, x::Number) = filter.b(x) ./ filter.a(x)
_freq(filter::ZeroPoleGain, x::Number) =
    filter.k * prod([x - z for z in filter.z]) / prod([x - p for p in filter.p])
_freq(filter::Biquad, x::Number) =
    ((filter.b0*x + filter.b1)*x + filter.b2) / ((x + filter.a1)*x  + filter.a2)
_freq(filter::SecondOrderSections, x::Number) =
    filter.g * prod([_freq(b, x) for b in filter.biquads])

"""
    phi, w = phaseresp(filter)

Phase response `phi` of a `filter` at (normalized) frequencies `w` in
radians/sample for a digital filter or radians/second for an analog filter
chosen as a reasonable default.
"""
function phaseresp(filter::FilterCoefficients)
    w = _freqrange(filter)
    return (phaseresp(filter, w), w)
end

"""
    phaseresp(filter, w)

Phase response of a `filter` at (normalized) frequency or frequencies `w` in
radians/sample for a digital filter or radians/second for an analog filter.
"""
function phaseresp(filter::FilterCoefficients, w)
    h = freqresp(filter, w)
    unwrap(angle.(h); dims=ndims(h))
end

"""
    tau, w = grpdelay(filter)

Group delay `tau` of a `filter` at (normalized) frequencies `w` in
radians/sample for a digital filter or radians/second for an analog filter
chosen as a reasonable default.
"""
function grpdelay(filter::FilterCoefficients)
    w = _freqrange(filter)
    return (grpdelay(filter, w), w)
end

"""
    grpdelay(fliter, w)

Group delay of a digital 'filter' at normalized frequency
or frequencies 'w' in radians/sample.
"""
function grpdelay(filter::FilterCoefficients{:z}, w)
    filter = convert(PolynomialRatio, filter)
    b, a = coefb(filter), coefa(filter)

    # Linear Phase FIR
    if (length(a) == 1) & (_is_sym(b) | _is_anti_sym(b))
        return fill((length(b)-1)/2, length(w))
    end

    c = xcorr(b, a; padmode = :none)
    cr = range(0, stop=length(c)-1) .* c
    ejw = exp.(-im .* w)
    num = Polynomial(cr).(ejw)
    den = Polynomial(c).(ejw)
    return real.(num ./ den) .- (length(a) - 1)
end

function grpdelay(filter::FilterCoefficients{:s}, w)
    filter = convert(PolynomialRatio, filter)
    b, a = filter.b, filter.a
    bd = derivative(b)
    ad = derivative(a)
    s = im .* w
    return real.((bd*a - ad*b).(s) ./ (a * b).(s))
end

"""
    impresp(filter, n=100)

Impulse response of a digital `filter` with `n` points.
"""
function impresp(filter::FilterCoefficients{:z}, n=100)
  i = [1; zeros(n-1)]
  filt(filter, i)
end

"""
    stepresp(filter, n=100)

Step response of a digital `filter` with `n` points.
"""
function stepresp(filter::FilterCoefficients{:z}, n=100)
  i = ones(n)
  filt(filter, i)
end

#
# Helper functions
#

function _is_sym(x::AbstractArray)
    n = length(x) ÷ 2
    return all(x[1+i] == x[end-i] for i in 0:n-1)
end

function _is_anti_sym(x::AbstractArray)
    n = length(x) ÷ 2
    return all(x[1+i] == -x[end-i] for i in 0:n)
end

_freqrange(::FilterCoefficients{:z}) = range(0, stop=π, length=257)
function _freqrange(filter::FilterCoefficients{:s})
    filter = convert(ZeroPoleGain, filter)
    w_interesting = sort!(Float64.(abs.([filter.p; filter.z])))
    include_zero = !isempty(w_interesting) && iszero(w_interesting[1])
    w_interesting = collect(Compat.Iterators.dropwhile(iszero, w_interesting))
    if isempty(w_interesting) # no non-zero poles or zeros
        if !include_zero || !isfinite(1/filter.k)
            return [0; 10 .^ (0:6)] # fallback
        end
        # include the point where |H|=1 (if any) and go further by factor 10
        return range(0.0, stop=10 * Float64(max(filter.k, 1/filter.k)), length=200)
    end
    # normal case: go from smalles to largest pole/zero, extended by factor 10
    w_min, w_max = w_interesting[[1,end]]
    w = 10 .^ range(log10(w_min)-1, stop=log10(w_max)+1, length=200)
    return include_zero ? [0.0; w] : w
end


include("_common.jl")

using GlobalSensitivity

function run_sa(X; node, metric)
    n_params = length(X)

    # drop dummy parameter
    X = X[1:n_params-1]
    update_params!(node, X...)

    Streamfall.run_node!(node, CALIB_CLIMATE)

    score = metric(CALIB_OBS[366:CALIB_LN], node.outflow[366:CALIB_LN])
    Streamfall.reset!(node)

    return score
end

N = 16384

for (app, metric) in zip(APPROACHES, OBJFUNCS)
    # sn, n_id = setup_network(joinpath(DATA_PATH, "cotter_baseline_calibrated_$(app).yml"))
    sn, n_id = setup_network(joinpath(DATA_PATH, "410730_IHACRES.yml"))
    node = sn[n_id]

    p_names, ini_vals, p_bounds = param_info(node; with_level=false)

    @info "Running $app..."

    # Set up dummy parameter
    push!(p_names, :dummy)
    push!(ini_vals, 0.0)
    push!(p_bounds, (0.0, 1000.0))

    sa_runner = (X) -> run_sa(X; node=node, metric=metric)
    sobol_results = gsa(sa_runner, Sobol(), p_bounds, N=N)

    headers = "$app Parameter, ST, S1, Threshold"
    res_lines = zip(p_names, sobol_results.ST, sobol_results.S1, repeat([0.04], length(p_names)))
    outfile = "$(DATA_PATH)sa_results_w_dummy.csv"

    open(outfile, "a+") do f
        println(f, headers)
        for line in res_lines
            println(f, join(line, ","))
        end

        print(f, "\n")
    end
end