export AbstractNLPModelMeta, NLPModelMeta, reset_data!

"""
    AbstractNLPModelMeta

Base type for metadata related to an optimization model.
"""
abstract type AbstractNLPModelMeta{T, S} end

"""
    NLPModelMeta <: AbstractNLPModelMeta

A composite type that represents the main features of the optimization problem

    optimize    obj(x)
    subject to  lvar ≤    x    ≤ uvar
                lcon ≤ cons(x) ≤ ucon

where `x`        is an `nvar`-dimensional vector,
      `obj`      is the real-valued objective function,
      `cons`     is the vector-valued constraint function,
      `optimize` is either "minimize" or "maximize".

Here, `lvar`, `uvar`, `lcon` and `ucon` are vectors.
Some of their components may be infinite to indicate that the corresponding bound or general constraint is not present.

---

    NLPModelMeta(nvar::Integer; kwargs...)
    NLPModelMeta(meta::AbstractNLPModelMeta; kwargs...)

Create an `NLPModelMeta` with `nvar` variables.
Alternatively, create an `NLPModelMeta` copy from another `AbstractNLPModelMeta`.
The following keyword arguments are accepted:
- `x0`: initial guess
- `lvar`: vector of lower bounds
- `uvar`: vector of upper bounds
- `nlvb`: number of nonlinear variables in both objectives and constraints
- `nlvo`: number of nonlinear variables in objectives (includes nlvb)
- `nlvc`: number of nonlinear variables in constraints (includes nlvb)
- `ncon`: number of general constraints
- `y0`: initial Lagrange multipliers
- `lcon`: vector of constraint lower bounds
- `ucon`: vector of constraint upper bounds
- `nnzo`: number of nonzeros in the gradient
- `nnzj`: number of elements needed to store the nonzeros in the Jacobian
- `lin_nnzj`: number of elements needed to store the nonzeros in the Jacobian of linear constraints
- `nln_nnzj`: number of elements needed to store the nonzeros in the Jacobian of nonlinear constraints
- `nnzh`: number of elements needed to store the nonzeros in the Hessian of the Lagrangian
- `lin`: indices of linear constraints
- `minimize`: true if optimize == minimize
- `islp`: true if the problem is a linear program
- `name`: problem name
- `variable_bounds_analysis`: whether to compute the partition of variables into fixed, lower-bounded, upper-bounded, range-bounded, free, and trivially infeasible sets
- `constraint_bounds_analysis`: whether to compute the partition of constraints into equality, lower-bounded, upper-bounded, range-bounded, free, and trivially infeasible sets
- `sparse_jacobian`: indicates whether the Jacobian of the constraints is sparse
- `sparse_hessian`: indicates whether the Hessian of the Lagrangian is sparse
- `grad_available`: indicates whether the gradient of the objective is available
- `jac_available`: indicates whether the Jacobian of the constraints is available
- `hess_available`: indicates whether the Hessian of the Lagrangian is available
- `jprod_available`: indicates whether the Jacobian-vector product `J * v` is available
- `jtprod_available`: indicates whether the transpose Jacobian-vector product `J' * v` is available
- `hprod_available`: indicates whether the Hessian-vector product of the Lagrangian `H * v` is available

`NLPModelMeta` also contains the following attributes, which are computed from the variables above:
- `nvar`: number of variables
- `ifix`: indices of fixed variables
- `ilow`: indices of variables with lower bound only
- `iupp`: indices of variables with upper bound only
- `irng`: indices of variables with lower and upper bound (range)
- `ifree`: indices of free variables
- `iinf`: indices of visibly infeasible bounds
- `jfix`: indices of equality constraints
- `jlow`: indices of constraints of the form c(x) ≥ cl
- `jupp`: indices of constraints of the form c(x) ≤ cu
- `jrng`: indices of constraints of the form cl ≤ c(x) ≤ cu
- `jfree`: indices of "free" constraints (there shouldn't be any)
- `jinf`: indices of the visibly infeasible constraints
- `nlin`: number of linear constraints
- `nnln`: number of nonlinear general constraints
- `nln`: indices of nonlinear constraints

The attributes `ifix`, `ilow`, `iupp`, `irng`, `ifree`, and `iinf` are not computed if `variable_bounds_analysis` is set to `false`.
The attributes `jfix`, `jlow`, `jupp`, `jrng`, `jfree`, and `jinf` are not computed if `constraint_bounds_analysis` is set to `false`.
"""
struct NLPModelMeta{T, S} <: AbstractNLPModelMeta{T, S}
  nvar::Int
  x0::S
  lvar::S
  uvar::S

  ifix::Vector{Int}
  ilow::Vector{Int}
  iupp::Vector{Int}
  irng::Vector{Int}
  ifree::Vector{Int}
  iinf::Vector{Int}

  nlvb::Int
  nlvo::Int
  nlvc::Int

  ncon::Int
  y0::S
  lcon::S
  ucon::S

  jfix::Vector{Int}
  jlow::Vector{Int}
  jupp::Vector{Int}
  jrng::Vector{Int}
  jfree::Vector{Int}
  jinf::Vector{Int}

  nnzo::Int
  nnzj::Int
  lin_nnzj::Int
  nln_nnzj::Int
  nnzh::Int

  nlin::Int
  nnln::Int

  lin::Vector{Int}
  nln::Vector{Int}

  minimize::Bool
  islp::Bool
  name::String

  variable_bounds_analysis::Bool
  constraint_bounds_analysis::Bool

  sparse_jacobian::Bool
  sparse_hessian::Bool

  grad_available::Bool
  jac_available::Bool
  hess_available::Bool
  jprod_available::Bool
  jtprod_available::Bool
  hprod_available::Bool
end

function NLPModelMeta{T, S}(
  nvar::Int;
  x0::S = fill!(S(undef, nvar), zero(T)),
  lvar::S = fill!(S(undef, nvar), T(-Inf)),
  uvar::S = fill!(S(undef, nvar), T(Inf)),
  nlvb = nvar,
  nlvo = nvar,
  nlvc = nvar,
  ncon = 0,
  y0::S = fill!(S(undef, ncon), zero(T)),
  lcon::S = fill!(S(undef, ncon), T(-Inf)),
  ucon::S = fill!(S(undef, ncon), T(Inf)),
  nnzo = nvar,
  nnzj = nvar * ncon,
  lin_nnzj = 0,
  nln_nnzj = nnzj - lin_nnzj,
  nnzh = nvar * (nvar + 1) ÷ 2,
  lin = Int[],
  minimize::Bool = true,
  islp::Bool = false,
  name = "Generic",
  variable_bounds_analysis::Bool = true,
  constraint_bounds_analysis::Bool = true,
  sparse_jacobian::Bool = true,
  sparse_hessian::Bool = true,
  grad_available::Bool = true,
  jac_available::Bool = (ncon > 0),
  hess_available::Bool = true,
  jprod_available::Bool = (ncon > 0),
  jtprod_available::Bool = (ncon > 0),
  hprod_available::Bool = true,
) where {T, S}
  if (nvar < 1) || (ncon < 0) || (nnzj < 0) || (nnzh < 0)
    error("Nonsensical dimensions")
  end

  @lencheck nvar x0 lvar uvar
  @lencheck ncon y0 lcon ucon
  @rangecheck 1 ncon lin
  @assert nnzj == lin_nnzj + nln_nnzj

  if variable_bounds_analysis
    ifix = findall(lvar .== uvar)
    ilow = findall((lvar .> T(-Inf)) .& (uvar .== T(Inf)))
    iupp = findall((lvar .== T(-Inf)) .& (uvar .< T(Inf)))
    irng = findall((lvar .> T(-Inf)) .& (uvar .< T(Inf)) .& (lvar .< uvar))
    ifree = findall((lvar .== T(-Inf)) .& (uvar .== T(Inf)))
    iinf = findall(lvar .> uvar)
  else
    ifix = Int[]
    ilow = Int[]
    iupp = Int[]
    irng = Int[]
    ifree = Int[]
    iinf = Int[]
  end

  if (ncon > 0) && constraint_bounds_analysis
    jfix = findall(lcon .== ucon)
    jlow = findall((lcon .> T(-Inf)) .& (ucon .== T(Inf)))
    jupp = findall((lcon .== T(-Inf)) .& (ucon .< T(Inf)))
    jrng = findall((lcon .> T(-Inf)) .& (ucon .< T(Inf)) .& (lcon .< ucon))
    jfree = findall((lcon .== T(-Inf)) .& (ucon .== T(Inf)))
    jinf = findall(lcon .> ucon)
  else
    jfix = Int[]
    jlow = Int[]
    jupp = Int[]
    jrng = Int[]
    jfree = Int[]
    jinf = Int[]
  end

  nln = setdiff(1:ncon, lin)
  nlin = length(lin)
  nnln = length(nln)

  NLPModelMeta{T, S}(
    nvar,
    x0,
    lvar,
    uvar,
    ifix,
    ilow,
    iupp,
    irng,
    ifree,
    iinf,
    nlvb,
    nlvo,
    nlvc,
    ncon,
    y0,
    lcon,
    ucon,
    jfix,
    jlow,
    jupp,
    jrng,
    jfree,
    jinf,
    nnzo,
    nnzj,
    lin_nnzj,
    nln_nnzj,
    nnzh,
    nlin,
    nnln,
    lin,
    nln,
    minimize,
    islp,
    name,
    variable_bounds_analysis,
    constraint_bounds_analysis,
    sparse_jacobian,
    sparse_hessian,
    grad_available,
    jac_available,
    hess_available,
    jprod_available,
    jtprod_available,
    hprod_available,
  )
end

NLPModelMeta(nvar::Int; x0::S = zeros(nvar), kwargs...) where {S} =
  NLPModelMeta{eltype(S), S}(nvar, x0 = x0; kwargs...)

function NLPModelMeta(
  meta::AbstractNLPModelMeta{T, S};
  nvar::Int = get_nvar(meta),
  x0::S = get_x0(meta),
  lvar::S = get_lvar(meta),
  uvar::S = get_uvar(meta),
  nlvb = get_nlvb(meta),
  nlvo = get_nlvo(meta),
  nlvc = get_nlvc(meta),
  ncon = get_ncon(meta),
  y0::S = get_y0(meta),
  lcon::S = get_lcon(meta),
  ucon::S = get_ucon(meta),
  nnzo = get_nnzo(meta),
  nnzj = get_nnzj(meta),
  lin_nnzj = get_lin_nnzj(meta),
  nln_nnzj = get_nln_nnzj(meta),
  nnzh = get_nnzh(meta),
  lin = get_lin(meta),
  minimize::Bool = get_minimize(meta),
  islp::Bool = get_islp(meta),
  name = get_name(meta),
  variable_bounds_analysis::Bool = get_variable_bounds_analysis(meta),
  constraint_bounds_analysis::Bool = get_constraint_bounds_analysis(meta),
  sparse_jacobian::Bool = get_sparse_jacobian(meta),
  sparse_hessian::Bool = get_sparse_hessian(meta),
  grad_available::Bool = get_grad_available(meta),
  jac_available::Bool = get_jac_available(meta),
  hess_available::Bool = get_hess_available(meta),
  jprod_available::Bool = get_jprod_available(meta),
  jtprod_available::Bool = get_jtprod_available(meta),
  hprod_available::Bool = get_hprod_available(meta),
) where {T, S}
  NLPModelMeta{T, S}(
    nvar,
    x0 = x0,
    lvar = lvar,
    uvar = uvar,
    nlvb = nlvb,
    nlvo = nlvo,
    nlvc = nlvc,
    ncon = ncon,
    y0 = y0,
    lcon = lcon,
    ucon = ucon,
    nnzo = nnzo,
    nnzj = nnzj,
    lin_nnzj = lin_nnzj,
    nln_nnzj = nln_nnzj,
    nnzh = nnzh,
    lin = lin,
    minimize = minimize,
    islp = islp,
    name = name,
    variable_bounds_analysis = variable_bounds_analysis,
    constraint_bounds_analysis = constraint_bounds_analysis,
    sparse_jacobian = sparse_jacobian,
    sparse_hessian = sparse_hessian,
    grad_available = grad_available,
    jac_available = jac_available,
    hess_available = hess_available,
    jprod_available = jprod_available,
    jtprod_available = jtprod_available,
    hprod_available = hprod_available,
  )
end

"""
    reset_data!(nlp)

Reset model data if appropriate.
This method should be overloaded if a subtype of `AbstractNLPModel`
contains data that should be reset, such as a quasi-Newton linear
operator.
"""
function reset_data!(nlp::AbstractNLPModel)
  return nlp
end
