#  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
#  Copyright (C) 2020-2023, UChicago Argonne, LLC. All rights reserved.
#  Released under the modified BSD license. See COPYING.md for more details.

import ..H5File

using OrderedCollections
using SparseArrays
using Statistics
using TimerOutputs

function collect_gmi(
    mps_filename;
    optimizer,
    max_rounds = 10,
    max_cuts_per_round = 100,
    atol = 1e-4,
)
    reset_timer!()

    # Open HDF5 file
    h5_filename = replace(mps_filename, ".mps.gz" => ".h5")
    h5 = H5File(h5_filename)

    # Read optimal solution
    sol_opt_dict = Dict(
        zip(
            h5.get_array("static_var_names"),
            convert(Array{Float64}, h5.get_array("mip_var_values")),
        ),
    )

    # Read optimal value
    obj_mip = h5.get_scalar("mip_lower_bound")
    if obj_mip === nothing
        obj_mip = h5.get_scalar("mip_obj_value")
    end
    obj_lp = h5.get_scalar("lp_obj_value")
    h5.file.close()

    # Define relative MIP gap
    gap(v) = 100 * abs(obj_mip - v) / abs(v)

    # Initialize stats
    stats_obj = []
    stats_gap = []
    stats_ncuts = []
    stats_time_convert = 0
    stats_time_solve = 0
    stats_time_select = 0
    stats_time_tableau = 0
    stats_time_gmi = 0
    all_cuts = nothing

    # Read problem
    model = read_from_file(mps_filename)

    for round = 1:max_rounds
        @info "Round $(round)..."

        stats_time_convert = @elapsed begin
            # Extract problem data
            data = ProblemData(model)

            # Construct optimal solution vector (with correct variable sequence)
            sol_opt = [sol_opt_dict[n] for n in data.var_names]

            # Assert optimal solution is feasible for the original problem
            assert_leq(data.constr_lb, data.constr_lhs * sol_opt)
            assert_leq(data.constr_lhs * sol_opt, data.constr_ub)

            # Convert to standard form
            data_s, transforms = convert_to_standard_form(data)
            model_s = to_model(data_s)
            set_optimizer(model_s, optimizer)
            relax_integrality(model_s)

            # Convert optimal solution to standard form
            sol_opt_s = forward(transforms, sol_opt)

            # Assert converted solution is feasible for standard form problem
            assert_eq(data_s.constr_lhs * sol_opt_s, data_s.constr_lb)
        end

        # Optimize standard form
        optimize!(model_s)
        stats_time_solve += solve_time(model_s)
        obj = objective_value(model_s)

        if round == 1
            # Assert standard form problem has same value as original
            assert_eq(obj, obj_lp)
            push!(stats_obj, obj)
            push!(stats_gap, gap(obj))
            push!(stats_ncuts, 0)
        end
        if termination_status(model_s) != MOI.OPTIMAL
            return
        end

        # Select tableau rows
        basis = get_basis(model_s)
        sol_frac = get_x(model_s)
        stats_time_select += @elapsed begin
            selected_rows =
                select_gmi_rows(data_s, basis, sol_frac, max_rows = max_cuts_per_round)
        end

        # Compute selected tableau rows
        stats_time_tableau += @elapsed begin
            tableau = compute_tableau(data_s, basis, sol_frac, rows = selected_rows)

            # Assert tableau rows have been computed correctly
            assert_eq(tableau.lhs * sol_frac, tableau.rhs)
            assert_eq(tableau.lhs * sol_opt_s, tableau.rhs)
        end

        # Compute GMI cuts
        stats_time_gmi += @elapsed begin
            cuts_s = compute_gmi(data_s, tableau)

            # Assert cuts have been generated correctly
            assert_cuts_off(cuts_s, sol_frac)
            assert_does_not_cut_off(cuts_s, sol_opt_s)

            # Abort if no cuts are left
            if length(cuts_s.lb) == 0
                @info "No cuts generated. Stopping."
                break
            end
        end

        # Add GMI cuts to original problem
        cuts = backwards(transforms, cuts_s)
        assert_does_not_cut_off(cuts, sol_opt)
        constrs = add_constraint_set(model, cuts)

        # Optimize original form
        set_optimizer(model, optimizer)
        undo_relax = relax_integrality(model)
        optimize!(model)
        obj = objective_value(model)
        push!(stats_obj, obj)
        push!(stats_gap, gap(obj))

        # Store useful cuts; drop useless ones from the problem
        useful = [abs(shadow_price(c)) > atol for c in constrs]
        drop = findall(useful .== false)
        keep = findall(useful .== true)
        delete.(model, constrs[drop])
        if all_cuts === nothing
            all_cuts = cuts
        else
            all_cuts.lhs = [all_cuts.lhs; cuts.lhs[keep, :]]
            all_cuts.lb = [all_cuts.lb; cuts.lb[keep]]
            all_cuts.ub = [all_cuts.ub; cuts.ub[keep]]
        end
        push!(stats_ncuts, length(all_cuts.lb))

        undo_relax()
    end

    # Store cuts
    if all_cuts !== nothing
        @info "Storing $(length(all_cuts.ub)) GMI cuts..."
        h5 = H5File(h5_filename)
        h5.put_sparse("cuts_lhs", all_cuts.lhs)
        h5.put_array("cuts_lb", all_cuts.lb)
        h5.put_array("cuts_ub", all_cuts.ub)
        h5.file.close()
    end

    return OrderedDict(
        "instance" => mps_filename,
        "max_rounds" => max_rounds,
        "rounds" => length(stats_obj) - 1,
        "time_convert" => stats_time_convert,
        "time_solve" => stats_time_solve,
        "time_tableau" => stats_time_tableau,
        "time_gmi" => stats_time_gmi,
        "obj_mip" => obj_mip,
        "stats_obj" => stats_obj,
        "stats_gap" => stats_gap,
        "stats_ncuts" => stats_ncuts,
    )
end

function select_gmi_rows(data, basis, x; max_rows = 10, atol = 0.001)
    candidate_rows = [
        r for r = 1:length(basis.var_basic) if (
            (data.var_types[basis.var_basic[r]] != 'C') &&
            (frac(x[basis.var_basic[r]]) > atol) &&
            (frac2(x[basis.var_basic[r]]) > atol)
        )
    ]
    candidate_vals = frac.(x[basis.var_basic[candidate_rows]])
    score = abs.(candidate_vals .- 0.5)
    perm = sortperm(score)
    return [candidate_rows[perm[i]] for i = 1:min(length(perm), max_rows)]
end

# function compute_gmi(data::ProblemData, tableau::Tableau)::ConstraintSet
#     @timeit "Initialization" begin
#         nrows, ncols = size(tableau.lhs)
#         ub = Float64[Inf for _ = 1:nrows]
#         lb = Float64[0.999 for _ = 1:nrows]
#         tableau_I, tableau_J, tableau_V = findnz(tableau.lhs)
#         lhs_I = Int[]
#         lhs_J = Int[]
#         lhs_V = Float64[]
#     end
#     @timeit "Compute coefficients" begin
#         for k = 1:nnz(tableau.lhs)
#             i::Int = tableau_I[k]
#             j::Int = tableau_J[k]
#             v::Float64 = 0.0
#             frac_alpha_j = frac(tableau_V[k])
#             alpha_j = tableau_V[k]
#             beta = frac(tableau.rhs[i])
#             if data.var_types[j] == 'C'
#                 if alpha_j >= 0
#                     v = alpha_j / beta
#                 else
#                     v = -alpha_j / (1 - beta)
#                 end
#             else
#                 if frac_alpha_j < beta
#                     v = frac_alpha_j / beta
#                 else
#                     v = (1 - frac_alpha_j) / (1 - beta)
#                 end
#             end
#             if abs(v) > 1e-8
#                 push!(lhs_I, i)
#                 push!(lhs_J, tableau_J[k])
#                 push!(lhs_V, v)
#             end
#         end
#     end
#     @timeit "Convert to ConstraintSet" begin
#         lhs = sparse(lhs_I, lhs_J, lhs_V, nrows, ncols)
#         cs = ConstraintSet(; lhs, ub, lb)
#     end
#     return cs
# end

function compute_gmi(data::ProblemData, tableau::Tableau)::ConstraintSet
    @timeit "Initialization" begin
        nrows::Int, ncols::Int = size(tableau.lhs)
        var_types::Vector{Char} = data.var_types
        tableau_rhs::Vector{Float64} = tableau.rhs
        tableau_I::Vector{Int}, tableau_J::Vector{Int}, tableau_V::Vector{Float64} = findnz(tableau.lhs)
    end

    @timeit "Pre-allocation" begin
        ub::Vector{Float64} = fill(Inf, nrows)
        lb::Vector{Float64} = fill(0.999, nrows)
        nnz_tableau::Int = length(tableau_I)
        lhs_I::Vector{Int} = Vector{Int}(undef, nnz_tableau)
        lhs_J::Vector{Int} = Vector{Int}(undef, nnz_tableau)
        lhs_V::Vector{Float64} = Vector{Float64}(undef, nnz_tableau)
        nnz_count::Int = 0
    end
    
    @timeit "Compute coefficients" begin
        @inbounds for k = 1:nnz_tableau
            i::Int = tableau_I[k]
            j::Int = tableau_J[k]
            alpha_j::Float64 = tableau_V[k]
            frac_alpha_j::Float64 = alpha_j - floor(alpha_j)
            beta_i::Float64 = tableau_rhs[i]
            beta::Float64 = beta_i - floor(beta_i)
            v::Float64 = 0
            
            # Compute coefficient
            if var_types[j] == 'C'
                if alpha_j >= 0
                    v = alpha_j / beta
                else
                    v = -alpha_j / (1 - beta)
                end
            else
                if frac_alpha_j < beta
                    v = frac_alpha_j / beta
                else
                    v = (1 - frac_alpha_j) / (1 - beta)
                end
            end
            
            # Store if significant
            if abs(v) > 1e-8
                nnz_count += 1
                lhs_I[nnz_count] = i
                lhs_J[nnz_count] = j
                lhs_V[nnz_count] = v
            end
        end
    end

    @timeit "Resize arrays to actual size" begin
        resize!(lhs_I, nnz_count)
        resize!(lhs_J, nnz_count)
        resize!(lhs_V, nnz_count)
    end

    @timeit "Convert to ConstraintSet" begin
        lhs::SparseMatrixCSC = sparse(lhs_I, lhs_J, lhs_V, nrows, ncols)
        cs::ConstraintSet = ConstraintSet(; lhs, ub, lb)
    end
    
    return cs
end


export compute_gmi,
    frac, select_gmi_rows, assert_cuts_off, assert_does_not_cut_off, collect_gmi