# Copyright (c) 2017: Miles Lubin and contributors
# Copyright (c) 2017: Google Inc.
#
# Use of this source code is governed by an MIT-style license that can be found
# in the LICENSE.md file or at https://opensource.org/licenses/MIT.

module TestSparseMatrix

import SparseArrays
using Test

import MathOptInterface as MOI

function runtests()
    for name in names(@__MODULE__; all = true)
        if startswith("$(name)", "test_")
            @testset "$(name)" begin
                getfield(@__MODULE__, name)()
            end
        end
    end
    return
end

"""
    test_empty_size()

Julia 1.7 is going to outlaw constructing empty sparse matrices. Test that we
can still do so.
"""
function test_empty_size()
    A = MOI.Utilities.MutableSparseMatrixCSC{
        Float64,
        Int,
        MOI.Utilities.ZeroBasedIndexing,
    }()
    MOI.empty!(A)
    @test A.rowval == Int[]
    @test A.nzval == Float64[]
    @test A.colptr == Int[0]
end

function test_empty_large()
    A = MOI.Utilities.MutableSparseMatrixCSC{
        Float64,
        Int,
        MOI.Utilities.ZeroBasedIndexing,
    }()
    MOI.empty!(A)
    MOI.Utilities.add_column(A)
    MOI.Utilities.add_column(A)
    MOI.Utilities.set_number_of_rows(A, 5)
    B = convert(SparseArrays.SparseMatrixCSC{Float64,Int}, A)
    @test size(B) == (5, 2)
end

function test_VectorAffine_ZeroBased()
    A = MOI.Utilities.MutableSparseMatrixCSC{
        Float64,
        Int,
        MOI.Utilities.ZeroBasedIndexing,
    }()
    MOI.empty!(A)
    x = MOI.VariableIndex.(1:3)
    f = MOI.VectorAffineFunction(
        vcat(
            MOI.VectorAffineTerm.(1, MOI.ScalarAffineTerm.(1.0, x)),
            MOI.VectorAffineTerm.(2, MOI.ScalarAffineTerm.([2.0, 3.0], x[2:3])),
        ),
        [0.5, 1.2],
    )

    index_map = MOI.Utilities.IndexMap()
    for i in 1:3
        MOI.Utilities.add_column(A)
        index_map[x[i]] = x[i]
    end
    MOI.Utilities.allocate_terms(A, index_map, f)
    MOI.Utilities.set_number_of_rows(A, 2)
    MOI.Utilities.load_terms(A, index_map, f, 0)
    MOI.Utilities.final_touch(A)
    B = convert(SparseArrays.SparseMatrixCSC{Float64,Int}, A)
    @test B == [1.0 1.0 1.0; 0.0 2.0 3.0]
    @test A.rowval == [0, 0, 1, 0, 1]
    @test A.nzval == [1.0, 1.0, 2.0, 1.0, 3.0]
    @test A.colptr == [0, 1, 3, 5]
    MOI.empty!(A)
    MOI.Utilities.add_column(A)
    C = convert(SparseArrays.SparseMatrixCSC{Float64,Int}, A)
    @test size(C) == (0, 1)
    return
end

function test_extract_function()
    A = MOI.Utilities.MutableSparseMatrixCSC{
        Float64,
        Int,
        MOI.Utilities.ZeroBasedIndexing,
    }()
    MOI.empty!(A)
    x = MOI.VariableIndex.(1:3)
    f = MOI.VectorAffineFunction(
        vcat(
            MOI.VectorAffineTerm.(1, MOI.ScalarAffineTerm.(1.0, x)),
            MOI.VectorAffineTerm.(2, MOI.ScalarAffineTerm.([2.0, 3.0], x[2:3])),
        ),
        [0.5, 1.2],
    )

    index_map = MOI.Utilities.IndexMap()
    MOI.Utilities.add_columns(A, 3)
    for i in 1:3
        index_map[x[i]] = x[i]
    end
    MOI.Utilities.allocate_terms(A, index_map, f)
    MOI.Utilities.set_number_of_rows(A, 2)
    MOI.Utilities.load_terms(A, index_map, f, 0)
    MOI.Utilities.final_touch(A)
    row_1 = MOI.Utilities.extract_function(A, 1, 0.5)
    @test row_1 ≈ MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(1.0, x), 0.5)
    @test MOI.Utilities.extract_function(A, 1:2, [0.5, 1.2]) ≈ f
    empty_f = MOI.Utilities.extract_function(A, 1:0, Float64[])
    @test empty_f ≈
          MOI.VectorAffineFunction(MOI.VectorAffineTerm{Float64}[], Float64[])
    r10 = MOI.Utilities.extract_function(A, 10, 0.5)
    @test r10 ≈ MOI.ScalarAffineFunction(MOI.ScalarAffineTerm{Float64}[], 0.5)
    f1 = 0.6 + 1.0 * x[1] + 1.0 * x[2] + 1.0 * x[3]
    r10 = MOI.Utilities.extract_function(A, 0:1, [0.5, 0.6])
    @test r10 ≈ MOI.Utilities.vectorize([0.5, f1])
    return
end

function test_VectorAffine_OneBased()
    A = MOI.Utilities.MutableSparseMatrixCSC{
        Float64,
        Int,
        MOI.Utilities.OneBasedIndexing,
    }()
    MOI.empty!(A)
    x = MOI.VariableIndex.(1:3)
    f = MOI.VectorAffineFunction(
        vcat(
            MOI.VectorAffineTerm.(1, MOI.ScalarAffineTerm.(1.0, x)),
            MOI.VectorAffineTerm.(2, MOI.ScalarAffineTerm.([2.0, 3.0], x[2:3])),
        ),
        [0.5, 1.2],
    )

    index_map = MOI.Utilities.IndexMap()
    for i in 1:3
        MOI.Utilities.add_column(A)
        index_map[x[i]] = x[i]
    end
    MOI.Utilities.allocate_terms(A, index_map, f)
    MOI.Utilities.set_number_of_rows(A, 2)
    MOI.Utilities.load_terms(A, index_map, f, 0)
    MOI.Utilities.final_touch(A)
    B = convert(SparseArrays.SparseMatrixCSC{Float64,Int}, A)
    @test B == [1.0 1.0 1.0; 0.0 2.0 3.0]
    @test A.rowval == [1, 1, 2, 1, 2]
    @test A.nzval == [1.0, 1.0, 2.0, 1.0, 3.0]
    @test A.colptr == [1, 2, 4, 6]
    MOI.empty!(A)
    MOI.Utilities.add_column(A)
    C = convert(SparseArrays.SparseMatrixCSC{Float64,Int}, A)
    @test size(C) == (0, 1)
    return
end

function test_ScalarAffine_ZeroBased()
    A = MOI.Utilities.MutableSparseMatrixCSC{
        Float64,
        Int,
        MOI.Utilities.ZeroBasedIndexing,
    }()
    MOI.empty!(A)
    x = MOI.VariableIndex.(1:3)
    f1 = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(1.0, x), 0.5)
    f2 =
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([2.0, 3.0], x[2:3]), 1.2)
    index_map = MOI.Utilities.IndexMap()
    for i in 1:3
        MOI.Utilities.add_column(A)
        index_map[x[i]] = x[i]
    end
    MOI.Utilities.allocate_terms(A, index_map, f1)
    MOI.Utilities.allocate_terms(A, index_map, f2)
    MOI.Utilities.set_number_of_rows(A, 2)
    MOI.Utilities.load_terms(A, index_map, f1, 0)
    MOI.Utilities.load_terms(A, index_map, f2, 1)
    MOI.Utilities.final_touch(A)
    B = convert(SparseArrays.SparseMatrixCSC{Float64,Int}, A)
    @test B == [1.0 1.0 1.0; 0.0 2.0 3.0]
    @test A.rowval == [0, 0, 1, 0, 1]
    @test A.nzval == [1.0, 1.0, 2.0, 1.0, 3.0]
    @test A.colptr == [0, 1, 3, 5]
    MOI.empty!(A)
    MOI.Utilities.add_column(A)
    C = convert(SparseArrays.SparseMatrixCSC{Float64,Int}, A)
    @test size(C) == (0, 1)
    return
end

function test_ScalarAffine_OneBased()
    A = MOI.Utilities.MutableSparseMatrixCSC{
        Float64,
        Int,
        MOI.Utilities.OneBasedIndexing,
    }()
    MOI.empty!(A)
    x = MOI.VariableIndex.(1:3)
    f1 = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(1.0, x), 0.5)
    f2 =
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([2.0, 3.0], x[2:3]), 1.2)
    index_map = MOI.Utilities.IndexMap()
    for i in 1:3
        MOI.Utilities.add_column(A)
        index_map[x[i]] = x[i]
    end
    MOI.Utilities.allocate_terms(A, index_map, f1)
    MOI.Utilities.allocate_terms(A, index_map, f2)
    MOI.Utilities.set_number_of_rows(A, 2)
    MOI.Utilities.load_terms(A, index_map, f1, 0)
    MOI.Utilities.load_terms(A, index_map, f2, 1)
    MOI.Utilities.final_touch(A)
    B = convert(SparseArrays.SparseMatrixCSC{Float64,Int}, A)
    @test B == [1.0 1.0 1.0; 0.0 2.0 3.0]
    @test A.rowval == [1, 1, 2, 1, 2]
    @test A.nzval == [1.0, 1.0, 2.0, 1.0, 3.0]
    @test A.colptr == [1, 2, 4, 6]
    MOI.empty!(A)
    MOI.Utilities.add_column(A)
    C = convert(SparseArrays.SparseMatrixCSC{Float64,Int}, A)
    @test size(C) == (0, 1)
    return
end

end

TestSparseMatrix.runtests()
