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

module TestLPSensitivity

using Test
using JuMP
import LinearAlgebra

struct TestSensitivitySolution
    primal::Float64
    dual::Float64
    basis::MOI.BasisStatusCode
    range::Tuple{Float64,Float64}
end

function _test_sensitivity(model_string, solution)
    m = MOIU.MockOptimizer(
        MOIU.Model{Float64}();
        eval_variable_constraint_dual = false,
    )
    model = direct_model(m)
    MOI.Utilities.loadfromstring!(m, model_string)
    optimize!(model)
    MOI.set(m, MOI.TerminationStatus(), MOI.OPTIMAL)
    MOI.set(m, MOI.PrimalStatus(), MOI.FEASIBLE_POINT)
    MOI.set(m, MOI.DualStatus(), MOI.FEASIBLE_POINT)
    obj_map = Dict{Any,Any}()
    for (key, val) in solution
        if key isa String
            var = variable_by_name(model, key)
            if var !== nothing
                obj_map[key] = var
                MOI.set(m, MOI.VariablePrimal(), index(var), val.primal)
                MOI.set(m, MOI.VariableBasisStatus(), index(var), val.basis)
            else
                c = constraint_by_name(model, key)
                @assert c !== nothing
                obj_map[key] = c
                MOI.set(m, MOI.ConstraintDual(), index(c), val.dual)
                MOI.set(m, MOI.ConstraintBasisStatus(), index(c), val.basis)
            end
        else
            var = variable_by_name(model, key[1])
            c = if key[2] <: MOI.GreaterThan
                LowerBoundRef(var)
            elseif key[2] <: MOI.LessThan
                UpperBoundRef(var)
            else
                FixRef(var)
            end
            @assert c !== nothing
            obj_map[key] = c
            MOI.set(m, MOI.ConstraintDual(), index(c), val.dual)
            MOI.set(m, MOI.ConstraintBasisStatus(), index(c), val.basis)
        end
    end
    sens = lp_sensitivity_report(model)
    for (s_key, val) in solution
        key = obj_map[s_key]
        @test sens[key][1] ≈ val.range[1]
        @test sens[key][2] ≈ val.range[2]
    end
    return
end

function test_Error_handling()
    m = MOIU.MockOptimizer(
        MOIU.Model{Float64}();
        eval_variable_constraint_dual = false,
    )
    model = direct_model(m)
    optimize!(model)
    MOI.set(m, MOI.TerminationStatus(), MOI.OPTIMAL)
    MOI.set(m, MOI.PrimalStatus(), MOI.NO_SOLUTION)
    MOI.set(m, MOI.DualStatus(), MOI.NO_SOLUTION)
    err = ErrorException(
        """
        Unable to compute LP sensitivity because no primal solution is available.

        Call `optimize!(model)` and verify `has_values(model)` returns `true` before calling `lp_sensitivity_report`.
        """,
    )
    @test_throws err lp_sensitivity_report(model)
    MOI.set(m, MOI.PrimalStatus(), MOI.FEASIBLE_POINT)
    err = ErrorException(
        """
        Unable to compute LP sensitivity because no dual solution is available.

        Call `optimize!(model)` and verify `has_duals(model)` returns `true` before calling `lp_sensitivity_report`.
        """,
    )
    @test_throws err lp_sensitivity_report(model)
    MOI.Utilities.loadfromstring!(
        m,
        """
variables: x
maxobjective: 1.0 * x
c: x in Interval(0.0, 1.0)
""",
    )
    err = ErrorException(
        """
        Unable to compute LP sensitivity because the model is not a linear program, or it contains interval constraints.

        Remove any nonlinear or integer components, and replace any interval constraints with separate upper and lower bound constraints.
        """,
    )
    @test_throws err lp_sensitivity_report(model)
    MOI.empty!(m)
    MOI.Utilities.loadfromstring!(
        m,
        """
variables: x, y
minobjective: 1.0 * x + 1.0 * y
c: [x, y] in Nonnegatives(2)
""",
    )
    err = ErrorException(
        """
        Unable to compute LP sensitivity because the model is not a linear program, or it contains interval constraints.

        Remove any nonlinear or integer components, and replace any interval constraints with separate upper and lower bound constraints.
        """,
    )
    @test_throws err lp_sensitivity_report(model)
    return
end

function test_Degeneracy()
    m = MOIU.MockOptimizer(
        MOIU.Model{Float64}();
        eval_variable_constraint_dual = false,
    )
    model = direct_model(m)
    MOI.Utilities.loadfromstring!(
        m,
        """
        variables: x, y
        maxobjective: 2.0 * x + 4.0 * y
        c1: x + 2 * y <= 4.0
        c2: x + 2 * y <= 4.0
        """,
    )
    optimize!(model)
    MOI.set(m, MOI.TerminationStatus(), MOI.OPTIMAL)
    MOI.set(m, MOI.PrimalStatus(), MOI.FEASIBLE_POINT)
    MOI.set(m, MOI.DualStatus(), MOI.FEASIBLE_POINT)
    for (key, val) in Dict("x" => 0, "y" => 2, "c1" => 0, "c2" => 2)
        var = variable_by_name(model, key)
        if var !== nothing
            MOI.set(m, MOI.VariablePrimal(), index(var), val)
            MOI.set(m, MOI.VariableBasisStatus(), index(var), MOI.BASIC)
        else
            c = constraint_by_name(model, key)
            MOI.set(m, MOI.ConstraintDual(), index(c), val)
            MOI.set(
                m,
                MOI.ConstraintBasisStatus(),
                index(c),
                MOI.NONBASIC_AT_UPPER,
            )
        end
    end
    @test_throws(LinearAlgebra.SingularException, lp_sensitivity_report(model))
    return
end

function test_Simple_bounds_Min()
    _test_sensitivity(
        """
        variables: x
        minobjective: x
        x >= -1.0
        x <= 1.0
        """,
        Dict(
            "x" => TestSensitivitySolution(
                -1.0,
                NaN,
                MOI.NONBASIC_AT_LOWER,
                (-1, Inf),
            ),
            ("x", MOI.GreaterThan{Float64}) =>
                TestSensitivitySolution(NaN, 1.0, MOI.NONBASIC, (-Inf, 2.0)),
            ("x", MOI.LessThan{Float64}) =>
                TestSensitivitySolution(NaN, 0.0, MOI.BASIC, (-2, Inf)),
        ),
    )
    return
end

function test_Simple_bounds_Max()
    _test_sensitivity(
        """
        variables: x
        maxobjective: x
        x >= -1.0
        x <= 1.0
        """,
        Dict(
            "x" => TestSensitivitySolution(
                1.0,
                NaN,
                MOI.NONBASIC_AT_UPPER,
                (-1, Inf),
            ),
            ("x", MOI.GreaterThan{Float64}) =>
                TestSensitivitySolution(NaN, 0.0, MOI.BASIC, (-Inf, 2.0)),
            ("x", MOI.LessThan{Float64}) =>
                TestSensitivitySolution(NaN, -1.0, MOI.NONBASIC, (-2, Inf)),
        ),
    )
    return
end

function test_Max_I()
    _test_sensitivity(
        """
        variables: x, y, z, w
        maxobjective: 1.1 * x + y
        x >= -1.0
        x <= 1.0
        y >= 0.0
        z == 1.0
        c1: x + y + z + w == 1.0
        c2: x + y         <= 2.0
        """,
        Dict(
            "x" => TestSensitivitySolution(
                1.0,
                NaN,
                MOI.NONBASIC_AT_UPPER,
                (-0.1, Inf),
            ),
            "y" => TestSensitivitySolution(1.0, NaN, MOI.BASIC, (-1, 0.1)),
            "z" => TestSensitivitySolution(1.0, NaN, MOI.NONBASIC, (-Inf, Inf)),
            "w" => TestSensitivitySolution(-2.0, NaN, MOI.BASIC, (-Inf, 1.0)),
            ("x", MOI.GreaterThan{Float64}) =>
                TestSensitivitySolution(NaN, 0.0, MOI.BASIC, (-Inf, 2.0)),
            ("x", MOI.LessThan{Float64}) =>
                TestSensitivitySolution(NaN, -0.1, MOI.NONBASIC, (-2, 1.0)),
            ("y", MOI.GreaterThan{Float64}) =>
                TestSensitivitySolution(NaN, 0.0, MOI.BASIC, (-Inf, 1.0)),
            "c1" =>
                TestSensitivitySolution(NaN, 0.0, MOI.NONBASIC, (-Inf, Inf)),
            "c2" =>
                TestSensitivitySolution(NaN, -1.0, MOI.NONBASIC, (-1.0, Inf)),
            ("z", MOI.EqualTo{Float64}) =>
                TestSensitivitySolution(NaN, 0.0, MOI.NONBASIC, (-Inf, Inf)),
        ),
    )
    return
end

function test_Min_I()
    _test_sensitivity(
        """
        variables: x, y, z, w
        minobjective: -1.1 * x + -1.0 * y
        x >= -1.0
        x <= 1.0
        y >= 0.0
        z == 1.0
        c1: x + y + z + w == 1.0
        c2: x + y         <= 2.0
        """,
        Dict(
            "x" => TestSensitivitySolution(
                1.0,
                NaN,
                MOI.NONBASIC_AT_UPPER,
                (-Inf, 0.1),
            ),
            "y" => TestSensitivitySolution(1.0, NaN, MOI.BASIC, (-0.1, 1)),
            "z" => TestSensitivitySolution(1.0, NaN, MOI.NONBASIC, (-Inf, Inf)),
            "w" => TestSensitivitySolution(-2.0, NaN, MOI.BASIC, (-1.0, Inf)),
            ("x", MOI.GreaterThan{Float64}) =>
                TestSensitivitySolution(NaN, 0.0, MOI.BASIC, (-Inf, 2.0)),
            ("x", MOI.LessThan{Float64}) =>
                TestSensitivitySolution(NaN, -0.1, MOI.NONBASIC, (-2, 1.0)),
            ("y", MOI.GreaterThan{Float64}) =>
                TestSensitivitySolution(NaN, 0.0, MOI.BASIC, (-Inf, 1.0)),
            "c1" =>
                TestSensitivitySolution(NaN, 0.0, MOI.NONBASIC, (-Inf, Inf)),
            "c2" =>
                TestSensitivitySolution(NaN, -1.0, MOI.NONBASIC, (-1.0, Inf)),
            ("z", MOI.EqualTo{Float64}) =>
                TestSensitivitySolution(NaN, 0.0, MOI.NONBASIC, (-Inf, Inf)),
        ),
    )
    return
end

function test_Max_II()
    _test_sensitivity(
        """
        variables: x, y
        maxobjective: -1.0 * x + -1.0 * y
        x >= 0.0
        y >= 0.0
        c1l: x + 2 * y >= -1.0
        c1u: x + 2 * y <= 2.0
        c2: x + y      >= 0.5
        c3: 2 * x + y  <= 2.0
        """,
        Dict(
            "x" => TestSensitivitySolution(0.5, NaN, MOI.BASIC, (0, 1)),
            "y" => TestSensitivitySolution(
                0.0,
                NaN,
                MOI.NONBASIC_AT_LOWER,
                (-Inf, 0),
            ),
            ("x", MOI.GreaterThan{Float64}) =>
                TestSensitivitySolution(NaN, 0.0, MOI.BASIC, (-Inf, 0.5)),
            ("y", MOI.GreaterThan{Float64}) =>
                TestSensitivitySolution(NaN, 0.0, MOI.NONBASIC, (-1, 0.5)),
            "c1l" => TestSensitivitySolution(NaN, 0.0, MOI.BASIC, (-Inf, 1.5)),
            "c1u" => TestSensitivitySolution(NaN, 0.0, MOI.BASIC, (-1.5, Inf)),
            "c2" =>
                TestSensitivitySolution(NaN, 1.0, MOI.NONBASIC, (-0.5, 0.5)),
            "c3" => TestSensitivitySolution(NaN, 1.0, MOI.BASIC, (-1.0, Inf)),
        ),
    )
    return
end

function test_Min_II()
    _test_sensitivity(
        """
        variables: x, y
        minobjective: 1.0 * x + 1.0 * y
        x >= 0.0
        y >= 0.0
        c1l: x + 2 * y >= -1.0
        c1u: x + 2 * y <= 2.0
        c2: x + y      >= 0.5
        c3: 2 * x + y  <= 2.0
        """,
        Dict(
            "x" => TestSensitivitySolution(0.5, NaN, MOI.BASIC, (-1, 0)),
            "y" => TestSensitivitySolution(
                0.0,
                NaN,
                MOI.NONBASIC_AT_LOWER,
                (0, Inf),
            ),
            ("x", MOI.GreaterThan{Float64}) =>
                TestSensitivitySolution(NaN, 0.0, MOI.BASIC, (-Inf, 0.5)),
            ("y", MOI.GreaterThan{Float64}) =>
                TestSensitivitySolution(NaN, 0.0, MOI.NONBASIC, (-1, 0.5)),
            "c1l" => TestSensitivitySolution(NaN, 0.0, MOI.BASIC, (-Inf, 1.5)),
            "c1u" => TestSensitivitySolution(NaN, 0.0, MOI.BASIC, (-1.5, Inf)),
            "c2" =>
                TestSensitivitySolution(NaN, 1.0, MOI.NONBASIC, (-0.5, 0.5)),
            "c3" => TestSensitivitySolution(NaN, 1.0, MOI.BASIC, (-1.0, Inf)),
        ),
    )
    return
end

function test_Max_III()
    _test_sensitivity(
        """
        variables: x, y
        maxobjective: 6.0 * x + 4.0 * y
        x >= 0.0
        y >= 0.0
        c1: 1 * x + 1 * y <=  6.0
        c2: 2 * x + 1 * y <=  9.0
        c3: 2 * x + 3 * y <= 16.0
        """,
        Dict(
            "x" => TestSensitivitySolution(3.0, NaN, MOI.BASIC, (-2, 2)),
            "y" => TestSensitivitySolution(3.0, NaN, MOI.BASIC, (-1, 2)),
            ("x", MOI.GreaterThan{Float64}) =>
                TestSensitivitySolution(NaN, 0.0, MOI.BASIC, (-Inf, 3)),
            ("y", MOI.GreaterThan{Float64}) =>
                TestSensitivitySolution(NaN, 0.0, MOI.BASIC, (-Inf, 3)),
            "c1" =>
                TestSensitivitySolution(NaN, -2.0, MOI.NONBASIC, (-1.5, 0.25)),
            "c2" => TestSensitivitySolution(NaN, -2.0, MOI.NONBASIC, (-1, 3)),
            "c3" => TestSensitivitySolution(NaN, 0.0, MOI.BASIC, (-1, Inf)),
        ),
    )
    return
end

function test_Min_III()
    _test_sensitivity(
        """
        variables: x, y
        minobjective: -6.0 * x + -4.0 * y
        x >= 0.0
        y >= 0.0
        c1: 1 * x + 1 * y <=  6.0
        c2: 2 * x + 1 * y <=  9.0
        c3: 2 * x + 3 * y <= 16.0
        """,
        Dict(
            "x" => TestSensitivitySolution(3.0, NaN, MOI.BASIC, (-2, 2)),
            "y" => TestSensitivitySolution(3.0, NaN, MOI.BASIC, (-2, 1)),
            ("x", MOI.GreaterThan{Float64}) =>
                TestSensitivitySolution(NaN, 0.0, MOI.BASIC, (-Inf, 3)),
            ("y", MOI.GreaterThan{Float64}) =>
                TestSensitivitySolution(NaN, 0.0, MOI.BASIC, (-Inf, 3)),
            "c1" =>
                TestSensitivitySolution(NaN, -2.0, MOI.NONBASIC, (-1.5, 0.25)),
            "c2" => TestSensitivitySolution(NaN, -2.0, MOI.NONBASIC, (-1, 3)),
            "c3" => TestSensitivitySolution(NaN, 0.0, MOI.BASIC, (-1, Inf)),
        ),
    )
    return
end

function test_Max_IV()
    _test_sensitivity(
        """
        variables: x, y
        maxobjective: 1.0 * x + 1.0 * y
        x >= 0.0
        y >= 0.0
        c1l: x + 2 * y >= -1.0
        c1u: x + 2 * y <= 2.0
        c2: x + y      >= 0.5
        c3: 2 * x + y  <= 2.0
        """,
        Dict(
            "x" => TestSensitivitySolution(2 / 3, NaN, MOI.BASIC, (-0.5, 1)),
            "y" => TestSensitivitySolution(2 / 3, NaN, MOI.BASIC, (-0.5, 1)),
            ("x", MOI.GreaterThan{Float64}) =>
                TestSensitivitySolution(NaN, 0.0, MOI.BASIC, (-Inf, 2 / 3)),
            ("y", MOI.GreaterThan{Float64}) =>
                TestSensitivitySolution(NaN, 0.0, MOI.BASIC, (-Inf, 2 / 3)),
            "c1l" => TestSensitivitySolution(NaN, 0.0, MOI.BASIC, (-Inf, 3)),
            "c1u" =>
                TestSensitivitySolution(NaN, -1 / 3, MOI.NONBASIC, (-1, 2)),
            "c2" => TestSensitivitySolution(NaN, 0.0, MOI.BASIC, (-Inf, 5 / 6)),
            "c3" =>
                TestSensitivitySolution(NaN, -1 / 3, MOI.NONBASIC, (-1.0, 2)),
        ),
    )
    return
end

function test_Min_IV()
    _test_sensitivity(
        """
        variables: x, y
        minobjective: -1.0 * x + -1.0 * y
        x >= 0.0
        y >= 0.0
        c1l: x + 2 * y >= -1.0
        c1u: x + 2 * y <= 2.0
        c2: x + y      >= 0.5
        c3: 2 * x + y  <= 2.0
        """,
        Dict(
            "x" => TestSensitivitySolution(2 / 3, NaN, MOI.BASIC, (-1, 0.5)),
            "y" => TestSensitivitySolution(2 / 3, NaN, MOI.BASIC, (-1, 0.5)),
            ("x", MOI.GreaterThan{Float64}) =>
                TestSensitivitySolution(NaN, 0.0, MOI.BASIC, (-Inf, 2 / 3)),
            ("y", MOI.GreaterThan{Float64}) =>
                TestSensitivitySolution(NaN, 0.0, MOI.BASIC, (-Inf, 2 / 3)),
            "c1l" => TestSensitivitySolution(NaN, 0.0, MOI.BASIC, (-Inf, 3)),
            "c1u" =>
                TestSensitivitySolution(NaN, -1 / 3, MOI.NONBASIC, (-1, 2)),
            "c2" => TestSensitivitySolution(NaN, 0.0, MOI.BASIC, (-Inf, 5 / 6)),
            "c3" =>
                TestSensitivitySolution(NaN, -1 / 3, MOI.NONBASIC, (-1.0, 2)),
        ),
    )
    return
end

function test_Free_variable()
    _test_sensitivity(
        """
        variables: x, y
        minobjective: 1.0 * x
        x >= 1.0
        """,
        Dict(
            "x" => TestSensitivitySolution(
                1.0,
                NaN,
                MOI.NONBASIC_AT_LOWER,
                (-1.0, Inf),
            ),
            "y" =>
                TestSensitivitySolution(0.0, NaN, MOI.SUPER_BASIC, (0.0, 0.0)),
            ("x", MOI.GreaterThan{Float64}) => TestSensitivitySolution(
                NaN,
                1.0,
                MOI.NONBASIC_AT_LOWER,
                (-Inf, Inf),
            ),
        ),
    )
    return
end

function test_fixed_to_set_basic_equality()
    _test_sensitivity(
        """
        variables: x, y
        maxobjective: 1.0 * x + 1.0 * y
        c: 1.0 * x + 1.0 * y <= 1.0
        x == 0.0
        """,
        Dict(
            "x" => TestSensitivitySolution(
                0.0,
                NaN,
                MOI.NONBASIC_AT_LOWER,
                (-Inf, Inf),
            ),
            "y" => TestSensitivitySolution(1.0, NaN, MOI.BASIC, (-1.0, Inf)),
            ("x", MOI.EqualTo{Float64}) => TestSensitivitySolution(
                NaN,
                0.0,
                MOI.NONBASIC_AT_LOWER,
                (0.0, 0.0),
            ),
            "c" =>
                TestSensitivitySolution(NaN, -1.0, MOI.NONBASIC, (-Inf, Inf)),
        ),
    )
    return
end

function test_variable_basis_status_error()
    inner = MOI.Utilities.MockOptimizer(MOI.Utilities.Model{Float64}())
    model = direct_model(inner)
    @variable(model, x >= 0)
    @constraint(model, c, x <= 1)
    optimize!(model)
    MOI.set(inner, MOI.TerminationStatus(), MOI.OPTIMAL)
    MOI.set(inner, MOI.PrimalStatus(), MOI.FEASIBLE_POINT)
    MOI.set(inner, MOI.DualStatus(), MOI.FEASIBLE_POINT)
    @test_throws(
        ErrorException(
            """
            Unable to compute LP sensitivity because the solver does not support querying the basis status of variables.

            Try a different LP solver that supports `MOI.VariableBasisStatus`.
            """,
        ),
        lp_sensitivity_report(model),
    )
    MOI.set(inner, MOI.VariableBasisStatus(), index(x), MOI.BASIC)
    @test_throws(
        ErrorException(
            """
            Unable to compute LP sensitivity because the solver does not support querying the basis status of constraints.

            Try a different LP solver that supports `MOI.ConstraintBasisStatus`.
            """,
        ),
        lp_sensitivity_report(model),
    )
    return
end

function test_degenerate_error()
    inner = MOI.Utilities.MockOptimizer(MOI.Utilities.Model{Float64}())
    model = direct_model(inner)
    @variable(model, x >= 0)
    optimize!(model)
    MOI.set(inner, MOI.TerminationStatus(), MOI.OPTIMAL)
    MOI.set(inner, MOI.PrimalStatus(), MOI.FEASIBLE_POINT)
    MOI.set(inner, MOI.DualStatus(), MOI.FEASIBLE_POINT)
    # If no error, this status should be NONBASIC_AT_LOWER...
    MOI.set(inner, MOI.VariableBasisStatus(), index(x), MOI.BASIC)
    @test_throws(
        ErrorException(
            """
            Unable to compute LP sensitivity because the basis matrix is not square, which indicates the problem is degenerate.

            Try adding bounds to free variables to obtain a non-degenerate basis.
            """,
        ),
        lp_sensitivity_report(model),
    )
    return
end

function test_is_lp()
    err = ErrorException(
        """
        Unable to compute LP sensitivity because the model is not a linear program, or it contains interval constraints.

        Remove any nonlinear or integer components, and replace any interval constraints with separate upper and lower bound constraints.
        """,
    )
    # Integer variable
    model = Model()
    @variable(model, x, Int)
    @test_throws err lp_sensitivity_report(model)
    # Quadratic objective
    model = Model()
    @variable(model, x)
    @objective(model, Min, x^2)
    @test_throws err lp_sensitivity_report(model)
    # Nonlinear objective
    model = Model()
    @variable(model, x)
    @objective(model, Max, log(x))
    @test_throws err lp_sensitivity_report(model)
    # Quadratic constraint
    model = Model()
    @variable(model, x)
    @constraint(model, x^2 <= 1)
    @test_throws err lp_sensitivity_report(model)
    # Conic constraint
    model = Model()
    @variable(model, x[1:3])
    @constraint(model, x in MOI.ExponentialCone())
    @test_throws err lp_sensitivity_report(model)
    return
end

end  # module
