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

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

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

function affs_are_equal(aff1::AffineExpression, aff2::AffineExpression)
    if aff1.constant != aff2.constant
        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 idx_to_coeff_1[v] != idx_to_coeff_2[v]
            return false
        end
    end
    return true
end

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

################################################################################
# PART A: Implement the following methods
# +(x::Real, y::Union{Variable, AffineExpression})
# +(x::Variable, y::Union{Real, Variable, AffineExpression})
# +(x::AffineExpression, y::Union{Real, Variable})
# +(x::AffineExpression, y::AffineExpression)
# -(x::Real, y::Union{Variable, AffineExpression})
# -(x::Variable, y::Union{Real, Variable, AffineExpression})
# -(x::AffineExpression, y::Union{Real, Variable})
# -(x::AffineExpression, y::AffineExpression)
# *(x::Real, y::Variable)
# *(x::Real, y::AffineExpression)
# *(x::Variable, y::Real)
# *(x::AffineExpression, y::Real)
# /(x::Union{Variable, AffineExpression}, y::Real)
#
# <=(lhs::Real, rhs::Union{Variable, AffineExpression})
# <=(lhs::Variable, rhs::Union{Real, Variable, AffineExpression})
# <=(lhs::AffineExpression, rhs::Union{Real, Variable, AffineExpression})
# ==(lhs::Real, rhs::Union{Variable, AffineExpression})
# ==(lhs::Variable, rhs::Union{Real, Variable, AffineExpression})
# ==(lhs::AffineExpression, rhs::Union{Real, Variable, AffineExpression})
# >=(lhs::Real, rhs::Union{Variable, AffineExpression})
# >=(lhs::Variable, rhs::Union{Real, Variable, AffineExpression})
# >=(lhs::AffineExpression, rhs::Union{Real, Variable, AffineExpression})
################################################################################


################################################################################
# PART B: Implement the following method
# construct_constraint_data(model::MyModel)::Tuple{SparseMatrixCSC{Float64,Int}, Vector{Float64}, Vector{Float64}}
################################################################################

end  # module OptModeling