# Simple min cost flow problem
#   
# n1 ------ f[1] -------> n3
#  \                     ^
#   \                   /
#    -f[2]-> n2 --f[3]--
# 
#  n1: demand = -10.1
#  n2: demand = 0
#  n3: demand = 10.1 
#  f[1]: cost = 0, capacity = 8.5, in mip model only integer flow allowed
#  f[2]: cost = 50, (capacity = 5 can be activated by removing comment at constraint, line 93)
#  f[3]: cost = 50
#
#  Correct solution for non-integer f[1]
#           f[1] = 8.5, f[2] = f[3] = 1.6, cost = 8.5*0 + 1.6*2*50 = 160
#  Correct solution for integer f[1]
#           f[1] = 8, f[2] = f[3] = 2.1, cost = 8.5*0 + 2.1*2*50 = 210
#

using JuMP
using BlockDecomposition
using Coluna
using GLPK

"""
colunaOptimizerForTreeSearch(useModifiedHeuristic)

Returns a coluna solver using the TreeSearchAlgorithm to solve the problem.

Parameters: useModifiedHeuristic: If true, uses SolveIpForm without enforcing integrality as heuristic,
                    if false, uses default (SolveIpForm with enforcing integrality).
"""
function colunaOptimizerForTreeSearch(useModifiedHeuristic)
    solver = Coluna.Algorithm.TreeSearchAlgorithm()
    if (useModifiedHeuristic)
        heuristic = Coluna.Algorithm.ParameterisedHeuristic(Coluna.Algorithm.SolveIpForm(get_dual_bound=false, enforce_integrality=false), 1.0, 1.0, 1, 1000, "Restricted Master IP")
        solver = Coluna.Algorithm.TreeSearchAlgorithm(
                        conqueralg=Coluna.Algorithm.ColCutGenConquer(
                            primal_heuristics=[heuristic]
                        )
                 )
    end
    return JuMP.optimizer_with_attributes(
        Coluna.Optimizer,
        "params" => Coluna.Params(
            solver=solver,
            global_art_var_cost=10e10,
            local_art_var_cost=10e10,
        ),
        "default_optimizer" => GLPK.Optimizer 
    );
end

"""
colunaOptimizerForColumnGeneration()

Returns a coluna solver using the ColumnGeneration directly to solve the problem.
"""
function colunaOptimizerForColumnGeneration()
    return JuMP.optimizer_with_attributes(
        Coluna.Optimizer,
        "params" => Coluna.Params(
            solver=Coluna.Algorithm.ColumnGeneration(),
            global_art_var_cost=10e10,
            local_art_var_cost=10e10,
        ),
        "default_optimizer" => GLPK.Optimizer 
    );
end

"""
solveFlowModel(f1isInteger, coluna)

Sets up the min-cost-flow example problem (see description at top of file) and solves it.

Parameters:
    f1isInteger: if true, flow variable for f[1] must be integer, if false, f[1] is continuous
        This means that if f1isInteger is true, the problem is a MIP, otherwise it is a purely continuous problem
    coluna: The coluna solver to setup the model with
"""
function solveFlowModel(f1isInteger, coluna)
    @BlockDecomposition.axis(M, 1:1)
    model = BlockDecomposition.BlockModel(coluna)
    @JuMP.variable(model, f[1:3, m in M])
    if f1isInteger
        JuMP.set_integer(f[1, 1])
    end
    @JuMP.constraint(model, n1[m in M], f[1,m] + f[2,m] == 10.1)
    @JuMP.constraint(model, n2[m in M], f[2,m] == f[3,m])
    @JuMP.constraint(model, n3[m in M], f[1,m] + f[3,m] == 10.1)
    @JuMP.constraint(model, cap1, sum(f[1,m] for m in M) <= 8.5)
    #@JuMP.constraint(model, cap2, sum(f[2,m] for m in M) <= 5)
    @JuMP.objective(model, Min, 50 * f[2,1] + 50 * f[3,1])

    @BlockDecomposition.dantzig_wolfe_decomposition(model, decomposition, M)

    subproblems = BlockDecomposition.getsubproblems(decomposition)
    BlockDecomposition.specify!.(subproblems, lower_multiplicity=1, upper_multiplicity=1)

    optimize!(model)

    println()
    println(" ------- Results -------------------")
    println()
    println("termination: $(termination_status(model))")
    println("primal_status: $(primal_status(model))")
    println()
    if f1isInteger
        println("cost = $(objective_value(model))   (expected: 210)")
        println("flows:")
        println("   f[1] = $(value(f[1,1]))  (expected: 8)")
        println("   f[2] = $(value(f[2,1]))  (expected: 2.1)")
        println("   f[3] = $(value(f[3,1]))  (expected: 2.1)")
    else
        println("cost = $(objective_value(model))   (expected: 160)")
        println("flows:")
        println("   f[1] = $(value(f[1,1]))  (expected: 8.5)")
        println("   f[2] = $(value(f[2,1]))  (expected: 1.6)")
        println("   f[3] = $(value(f[3,1]))  (expected: 1.6)")
    end

end