#  Copyright 2017, Iain Dunning, Joey Huchette, Miles Lubin, and contributors
#  This Source Code Form is subject to the terms of the Mozilla Public
#  License, v. 2.0. If a copy of the MPL was not distributed with this
#  file, You can obtain one at https://mozilla.org/MPL/2.0/.
#############################################################################
# JuMP
# An algebraic modeling language for Julia
# See https://github.com/jump-dev/JuMP.jl
#############################################################################

module TestMutableArithmetics

using JuMP
using Test

struct DummyVariableRef <: AbstractVariableRef end

JuMP.name(::DummyVariableRef) = "dummy"

struct DummyExpr end

function Base.:(+)(
    ::GenericAffExpr{Float64,DummyVariableRef},
    ::AbstractJuMPScalar,
)
    return DummyExpr()
end

function Base.:(+)(
    ::AbstractJuMPScalar,
    ::GenericAffExpr{Float64,DummyVariableRef},
)
    return DummyExpr()
end

function Base.:(-)(
    ::GenericAffExpr{Float64,DummyVariableRef},
    ::AbstractJuMPScalar,
)
    return DummyExpr()
end

function Base.:(-)(
    ::AbstractJuMPScalar,
    ::GenericAffExpr{Float64,DummyVariableRef},
)
    return DummyExpr()
end

function JuMP._MA.promote_operation(
    ::Union{typeof(+),typeof(-)},
    ::Type{DummyVariableRef},
    ::Type{<:AbstractVariableRef},
)
    return DummyExpr
end

function JuMP._MA.promote_operation(
    ::Union{typeof(+),typeof(-)},
    ::Type{<:AbstractVariableRef},
    ::Type{DummyVariableRef},
)
    return DummyExpr
end

function JuMP._MA.promote_operation(
    ::Union{typeof(+),typeof(-)},
    ::Type{DummyVariableRef},
    ::Type{<:AbstractJuMPScalar},
)
    return DummyExpr
end

function JuMP._MA.promote_operation(
    ::Union{typeof(+),typeof(-)},
    ::Type{<:AbstractJuMPScalar},
    ::Type{DummyVariableRef},
)
    return DummyExpr
end

function promote_operation_test(op::Function, x::Type, y::Type)
    f() = JuMP._MA.promote_operation(op, x, y)
    @test typeof(op(zero(x), zero(y))) == f()
    @test 0 == @allocated f()
    return
end

function sym_promote_operation_test(op::Function, x::Type, y::Type)
    promote_operation_test(op, x, y)
    promote_operation_test(op, y, x)
    return
end

function test_extension_promote_operation(
    ModelType = Model,
    VariableRefType = VariableRef,
)
    for S in [Float64, ComplexF64]
        AffExprType = GenericAffExpr{S,VariableRefType}
        QuadExprType = GenericQuadExpr{S,VariableRefType}
        for op in [+, -, *]
            for T in [Int, S]
                sym_promote_operation_test(op, T, VariableRefType)
                sym_promote_operation_test(op, T, AffExprType)
                sym_promote_operation_test(op, T, QuadExprType)
            end
            promote_operation_test(op, VariableRefType, VariableRefType)
            sym_promote_operation_test(op, VariableRefType, AffExprType)
            if op != *
                sym_promote_operation_test(op, VariableRefType, QuadExprType)
                sym_promote_operation_test(op, AffExprType, QuadExprType)
            end
        end
    end
    S = Float64
    T = ComplexF64
    for op in [+, -, *]
        AffS = GenericAffExpr{S,VariableRefType}
        AffT = GenericAffExpr{T,VariableRefType}
        QuadS = GenericQuadExpr{S,VariableRefType}
        QuadT = GenericQuadExpr{T,VariableRefType}
        sym_promote_operation_test(op, AffS, AffT)
        if op != *
            sym_promote_operation_test(op, AffS, QuadT)
            sym_promote_operation_test(op, AffT, QuadS)
            sym_promote_operation_test(op, QuadS, QuadT)
        end
    end
    return
end

function test_int()
    model = Model()
    @variable(model, x)
    for a in [1, 1im]
        JuMP._MA.Test.int_test(
            typeof(a * x);
            exclude = ["int_mul", "int_add", "int_add_mul"],
        )
        JuMP._MA.Test.int_test(
            typeof(a * x^2);
            exclude = ["int_mul", "int_add", "int_add_mul"],
        )
    end
    return
end

function test_scalar()
    model = Model()
    @variable(model, x)
    exclude = ["cube"]
    JuMP._MA.Test.scalar_test(x; exclude = exclude)
    for a in [2, 2im]
        JuMP._MA.Test.scalar_test(a * x + 3; exclude = exclude)
        JuMP._MA.Test.scalar_test(a * x^2 + 4x + 1; exclude = exclude)
    end
    return
end

function test_quadratic()
    model = Model()
    @variable(model, x[1:4])
    # Test is excluded because of https://github.com/jump-dev/MutableArithmetics.jl/issues/227
    JuMP._MA.Test.quadratic_test(x...; exclude = ["quadratic_add_canonical"])
    return
end

function test_sparse()
    model = Model()
    @variable(model, X11)
    @variable(model, X23)
    @variable(model, Xd[1:3, 1:3])
    JuMP._MA.Test.sparse_test(X11, X23, Xd)
    return
end

function test_vector()
    model = Model()
    @variable(model, x[1:3])
    JuMP._MA.Test.array_test(x)
    return
end

function test_symmetric()
    model = Model()
    @variable(model, y[1:2, 1:2], Symmetric)
    JuMP._MA.Test.array_test(y)
    return
end

function test_nonsquare_matrix()
    model = Model()
    @variable(model, z[1:2, 1:3])
    JuMP._MA.Test.array_test(z)
    return
end

function test_DenseAxisVector()
    model = Model()
    @variable(model, y[2:5])
    JuMP._MA.Test.array_test(y; exclude = ["matrix_vector", "non_array"])
    return
end

function test_different_variables()
    model = Model()
    x = @variable(model)
    y = DummyVariableRef()
    for a in (x, x + 1, x + im)
        @test JuMP._MA.promote_operation(+, typeof(a), typeof(y)) == DummyExpr
        @test JuMP._MA.promote_operation(-, typeof(a), typeof(y)) == DummyExpr
        @test JuMP._MA.promote_operation(+, typeof(y), typeof(a)) == DummyExpr
        @test JuMP._MA.promote_operation(-, typeof(y), typeof(a)) == DummyExpr
    end
    return
end

function test_scaling_errors()
    model = Model()
    @variable(model, x)
    @test_throws InexactError JuMP._MA.scaling(x + 1.0)
    @test JuMP._MA.scaling(zero(AffExpr)) == 0.0
    @test JuMP._MA.scaling(AffExpr(1.0)) == 1.0
    @test_throws InexactError JuMP._MA.scaling(x^2 + 1.0)
    @test JuMP._MA.scaling(zero(QuadExpr)) == 0.0
    @test JuMP._MA.scaling(QuadExpr(AffExpr(1.0))) == 1.0
    return
end

function test_operate_zero()
    model = Model()
    @variable(model, x)
    f = 1.0 * x + 2.0
    @test JuMP._MA.operate!(*, f, 0.0) == zero(AffExpr)
    return
end

end
