module OptModeling

using SparseArrays

export MyModel, Variable, AffineExpression, LinearConstraint, add_variable, add_constraint, construct_constraint_data

mutable struct MyModel
    num_variables::Int
    constraints::Vector
    MyModel() = new(0, Any[])
end

struct Variable
    model::MyModel
    index::Int
end
Base.show(io::IO, x::Variable) = Base.print(io, "Variable($(x.index))")

function add_variable(model::MyModel)
    model.num_variables += 1
    return Variable(model, model.num_variables)
end

struct AffineExpression
    variables::Vector{Variable}
    coeffs::Vector{Float64}
    constant::Float64
end
AffineExpression(x::Real) = AffineExpression(Variable[], Float64[], x)
AffineExpression(x::Variable) = AffineExpression([x], [1.0], 0.0)
AffineExpression(x::AffineExpression) = AffineExpression(copy(x.variables), copy(x.coeffs), x.constant)

struct LinearConstraint
    aff_expr::AffineExpression
    lb::Float64
    ub::Float64
end

add_constraint(model::MyModel, constr::LinearConstraint) = push!(model.constraints, constr)

TOL = 1e-8

function affs_are_equal(aff1::AffineExpression, aff2::AffineExpression)
    if !(abs(aff1.constant - aff2.constant) < TOL)
        return false
    end
    idx_to_coeff_1 = Dict{Int,Float64}()
    idx_to_coeff_2 = Dict{Int,Float64}()
    for i in 1:length(aff1.variables)
        idx = aff1.variables[i].index
        init_value = get(idx_to_coeff_1, idx, 0.0)
        idx_to_coeff_1[idx] = init_value + aff1.coeffs[i]
    end
    for i in 1:length(aff2.variables)
        idx = aff2.variables[i].index
        init_value = get(idx_to_coeff_2, idx, 0.0)
        idx_to_coeff_2[idx] = init_value + aff2.coeffs[i]
    end
    if keys(idx_to_coeff_1) != keys(idx_to_coeff_2)
        return false
    end
    for v in keys(idx_to_coeff_1)
        if !(abs(idx_to_coeff_1[v] - idx_to_coeff_2[v]) < TOL)
            return false
        end
    end
    return true
end

import Base: +, -, *, /, <=, ==, >=

################################################################################
# PART A: Implement the following methods
+(x::Real, y::Union{Variable, AffineExpression}) = AffineExpression(x) + AffineExpression(y)
+(x::Variable, y::Union{Real, Variable, AffineExpression}) = AffineExpression(x) + AffineExpression(y)
+(x::AffineExpression, y::Union{Real, Variable}) = x + AffineExpression(y)
+(x::AffineExpression, y::AffineExpression) = AffineExpression(vcat(x.variables, y.variables), vcat(x.coeffs, y.coeffs), x.constant + y.constant)
-(x::Real, y::Union{Variable, AffineExpression}) = AffineExpression(x) - AffineExpression(y)
-(x::Variable, y::Union{Real, Variable, AffineExpression}) = AffineExpression(x) - AffineExpression(y)
-(x::AffineExpression, y::Union{Real, Variable}) = x - AffineExpression(y)
-(x::AffineExpression, y::AffineExpression) = AffineExpression(vcat(x.variables, y.variables), vcat(x.coeffs, -y.coeffs), x.constant - y.constant)
*(x::Real, y::Variable) = AffineExpression([y], Float64[x], 0.0)
*(x::Real, y::AffineExpression) = AffineExpression(copy(y.variables), x .* y.coeffs, x * y.constant)
*(x::Variable, y::Real) = AffineExpression([x], Float64[y], 0.0)
*(x::AffineExpression, y::Real) = AffineExpression(copy(x.variables), y .* x.coeffs, y * x.constant)
/(x::Union{Variable, AffineExpression}, y::Real) = x * (1 / y)

<=(lhs::Real, rhs::Union{Variable, AffineExpression}) = LinearConstraint(lhs - rhs, -Inf, 0.0)
<=(lhs::Variable, rhs::Union{Real, Variable, AffineExpression}) = LinearConstraint(lhs - rhs, -Inf, 0.0)
<=(lhs::AffineExpression, rhs::Union{Real, Variable, AffineExpression}) = LinearConstraint(lhs - rhs, -Inf, 0.0)
==(lhs::Real, rhs::Union{Variable, AffineExpression}) = LinearConstraint(lhs - rhs, 0.0, 0.0)
==(lhs::Variable, rhs::Union{Real, Variable, AffineExpression}) = LinearConstraint(lhs - rhs, 0.0, 0.0)
==(lhs::AffineExpression, rhs::Union{Real, Variable, AffineExpression}) = LinearConstraint(lhs - rhs, 0.0, 0.0)
>=(lhs::Real, rhs::Union{Variable, AffineExpression}) = LinearConstraint(lhs - rhs, 0.0, Inf)
>=(lhs::Variable, rhs::Union{Real, Variable, AffineExpression}) = LinearConstraint(lhs - rhs, 0.0, Inf)
>=(lhs::AffineExpression, rhs::Union{Real, Variable, AffineExpression}) = LinearConstraint(lhs - rhs, 0.0, Inf)
################################################################################


################################################################################
# PART B: Implement the following method
function construct_constraint_data(model::MyModel)::Tuple{SparseMatrixCSC{Float64,Int}, Vector{Float64}, Vector{Float64}}
    I = Int[]
    J = Int[]
    V = Float64[]
    lb = Float64[]
    ub = Float64[]
    m = length(model.constraints)
    n = model.num_variables
    for (i,constr) in enumerate(model.constraints)
        aff = constr.aff_expr
        for idx in 1:length(aff.variables)
            j = aff.variables[idx].index
            val = aff.coeffs[idx]
            push!(I, i)
            push!(J, j)
            push!(V, val)
        end
        push!(lb, constr.lb - aff.constant)
        push!(ub, constr.ub - aff.constant)
    end
    return sparse(I, J, V, m, n), lb, ub
end
################################################################################

end  # module OptModeling