ommx 3.0.0-beta.2

use std::collections::{BTreeMap, HashMap, HashSet};

use super::{
    format::{CONSTR_PREFIX, OBJ_NAME, VAR_PREFIX},
    parser::{ColumnName, ObjSense, RowName},
    Mps,
};
use crate::{decision_variable::Kind as DecisionVariableKind, Coefficient, Quadratic};
use crate::{
    v1, Bound, Constraint, ConstraintContext, ConstraintID, DecisionVariable, Equality, Function,
    Instance, ModelingLabel, Sense, VariableID,
};

type ConvertedDecisionVariables = (
    BTreeMap<VariableID, DecisionVariable>,
    Vec<(VariableID, String)>,
    HashMap<ColumnName, VariableID>,
);

type ConvertedConstraints = (
    BTreeMap<ConstraintID, Constraint>,
    Vec<(ConstraintID, String)>,
);

pub fn convert(mps: Mps) -> crate::Result<Instance> {
    let (decision_variables, var_names, name_id_map) = convert_dvars(&mps);
    let objective = convert_objective(&mps, &name_id_map)?;
    let (constraints, constraint_names) = convert_constraints(&mps, &name_id_map)?;
    let sense = convert_sense(mps.obj_sense);

    let mut instance = Instance::new(sense, objective, decision_variables, constraints)?;

    // Drain MPS names as modeling labels through the instance owner boundary;
    // per-element label storage was retired in v3.
    for (id, name) in var_names {
        instance.set_variable_label(
            id,
            ModelingLabel {
                name: Some(name),
                ..Default::default()
            },
        )?;
    }
    for (id, name) in constraint_names {
        instance.set_constraint_context(
            id,
            ConstraintContext {
                label: ModelingLabel {
                    name: Some(name),
                    ..Default::default()
                },
                ..Default::default()
            },
        )?;
    }

    instance.description = convert_description(&mps);

    Ok(instance)
}

fn convert_description(mps: &Mps) -> Option<v1::instance::Description> {
    // currently only gets the name
    if mps.name.is_empty() {
        None
    } else {
        Some(v1::instance::Description {
            name: Some(mps.name.to_owned()),
            ..Default::default()
        })
    }
}

fn convert_dvars(mps: &Mps) -> ConvertedDecisionVariables {
    let Mps {
        vars,
        u,
        l,
        integer,
        binary,
        ..
    } = mps;
    let mut dvars = BTreeMap::new();
    let mut var_names: Vec<(VariableID, String)> = Vec::new();
    // Will be used to keep track of dvar ids throughout the conversion. might
    // not be strictly necessary if we make it so parser.rs and this file
    // guarantee a strict and consistent ordering, but it's less error-prone
    // this way
    let mut name_id_map = HashMap::with_capacity(u.len());

    // We want to be able to recover IDs if the var name is in the form
    // "OMMX_VAR_<number>", matching how our output formats names.
    //
    // NOTE considering the case where an OMMX-created MPS file was later
    // edited, there may be a mix of valid OMMX id names and invalid names.
    // Handling all edge cases would be pretty complex and potentially bad for
    // performance. For simplicity, we only apply ID recovery when ALL variables
    // match the naming pattern.
    if vars.iter().any(|name| !name.starts_with(VAR_PREFIX)) {
        // general case -- assign ids by order
        for (i, var_name) in vars.iter().enumerate() {
            let kind = get_dvar_kind(var_name, integer, binary);
            let bound = get_dvar_bound(var_name, l, u);
            // our ID ends up being dependent on the order of vars hashset. This is
            // unstable across executions -- we might want to consider an indexset
            // in the future
            let id = VariableID::from(i as u64);
            name_id_map.insert(var_name.clone(), id);
            let dvar = DecisionVariable::new(kind, bound, crate::ATol::default())
                .expect("Failed to create decision variable");
            dvars.insert(id, dvar);
            var_names.push((id, var_name.0.clone()));
        }
    } else {
        // recover IDs case
        for (id_value, var_name) in vars
            .iter()
            .filter_map(|name| parse_id_tag(VAR_PREFIX, name).map(|id| (id, name)))
        {
            let kind = get_dvar_kind(var_name, integer, binary);
            let bound = get_dvar_bound(var_name, l, u);
            let id = VariableID::from(id_value);
            name_id_map.insert(var_name.clone(), id);
            let dvar = DecisionVariable::new(kind, bound, crate::ATol::default())
                .expect("Failed to create decision variable");
            // Do not add name here since it means OMMX ID, not user-defined name
            dvars.insert(id, dvar);
        }
    }
    (dvars, var_names, name_id_map)
}

/// Strips the prefix of a variable/constraint name, and parses the following id number.
///
/// Returns none if prefix is not present, or if parsing as u64 fails.
fn parse_id_tag(prefix: &str, name: &str) -> Option<u64> {
    name.strip_prefix(prefix)?.parse().ok()
}

// name_id_map helps us convert from column name to id.
// See comment in `convert_dvars`
fn convert_objective(
    mps: &Mps,
    name_id_map: &HashMap<ColumnName, VariableID>,
) -> crate::Result<Function> {
    let Mps { b, c, quad_obj, .. } = mps;
    let constant = -b.get(&OBJ_NAME.into()).copied().unwrap_or_default();

    // Check if we have any quadratic terms
    if quad_obj.is_empty() {
        // Linear objective only
        if c.is_empty() {
            Ok(Function::try_from(constant)?)
        } else {
            // Build linear function by adding terms
            let mut linear = crate::Linear::try_from(constant)?;
            for (name, &coefficient) in c {
                if let Some(&id) = name_id_map.get(name) {
                    match Coefficient::try_from(coefficient) {
                        Ok(coef) => {
                            linear.add_term(crate::LinearMonomial::Variable(id), coef)?;
                        }
                        Err(crate::CoefficientError::Zero) => {} // Skip zero coefficients
                        Err(e) => return Err(e.into()),
                    }
                }
            }
            Ok(if linear.is_zero() {
                Function::Zero
            } else {
                Function::Linear(linear)
            })
        }
    } else {
        // Quadratic objective - need to build a quadratic function
        let mut quadratic = Quadratic::try_from(constant)?;

        // Add linear terms
        for (name, &coefficient) in c {
            if let Some(&id) = name_id_map.get(name) {
                match Coefficient::try_from(coefficient) {
                    Ok(coef) => {
                        quadratic.add_term(crate::QuadraticMonomial::Linear(id), coef)?;
                    }
                    Err(crate::CoefficientError::Zero) => {} // Skip zero coefficients
                    Err(e) => return Err(e.into()),
                }
            }
        }

        // Add quadratic terms
        for ((var1_name, var2_name), &coefficient) in quad_obj {
            if let (Some(&id1), Some(&id2)) =
                (name_id_map.get(var1_name), name_id_map.get(var2_name))
            {
                match Coefficient::try_from(coefficient) {
                    Ok(coef) => {
                        quadratic.add_term(crate::QuadraticMonomial::new_pair(id1, id2), coef)?;
                    }
                    Err(crate::CoefficientError::Zero) => {} // Skip zero coefficients
                    Err(e) => return Err(e.into()),
                }
            }
        }

        Ok(Function::Quadratic(quadratic))
    }
}

fn convert_constraints(
    mps: &Mps,
    name_id_map: &HashMap<ColumnName, VariableID>,
) -> crate::Result<ConvertedConstraints> {
    let Mps {
        a,
        b,
        eq,
        ge,
        le,
        quad_constraints,
        ..
    } = mps;
    let mut constrs = BTreeMap::new();
    let mut constraint_names: Vec<(ConstraintID, String)> = Vec::new();

    // as with decision variables, we're trying to recover IDs whenever all constraints match the naming scheme
    if a.keys().any(|name| !name.starts_with(CONSTR_PREFIX)) {
        // general case -- assign ids by order
        for (i, (row_name, row)) in a.iter().enumerate() {
            let b_value = b.get(row_name).copied().unwrap_or(0.0);
            let quad_terms = quad_constraints.get(row_name);
            let (function, equality) =
                convert_inequality(row, b_value, row_name, eq, ge, le, name_id_map, quad_terms)?;
            let id = ConstraintID::from(i as u64);
            let constraint = match equality {
                Equality::EqualToZero => Constraint::equal_to_zero(function),
                Equality::LessThanOrEqualToZero => Constraint::less_than_or_equal_to_zero(function),
            };
            constrs.insert(id, constraint);
            constraint_names.push((id, row_name.0.clone()));
        }
    } else {
        // recover IDs case
        for (id_value, row_name, row) in a.iter().filter_map(|(row_name, row)| {
            parse_id_tag(CONSTR_PREFIX, row_name).map(|id| (id, row_name, row))
        }) {
            let b_value = b.get(row_name).copied().unwrap_or(0.0);
            let quad_terms = quad_constraints.get(row_name);
            let (function, equality) =
                convert_inequality(row, b_value, row_name, eq, ge, le, name_id_map, quad_terms)?;
            let id = ConstraintID::from(id_value);
            let constraint = match equality {
                Equality::EqualToZero => Constraint::equal_to_zero(function),
                Equality::LessThanOrEqualToZero => Constraint::less_than_or_equal_to_zero(function),
            };
            // Do not add name here since it means OMMX ID, not user-defined name
            constrs.insert(id, constraint);
        }
    }
    Ok((constrs, constraint_names))
}

/// Handles passing the `b` constant part to the left-hand side, as we only
/// accept the right-hand side being 0.0.
///
/// Returns the full function plus what the OMMX equality should be.
#[allow(clippy::too_many_arguments)]
fn convert_inequality(
    row: &HashMap<ColumnName, f64>,
    mut b: f64,
    name: &RowName,
    eq: &HashSet<RowName>,
    ge: &HashSet<RowName>,
    le: &HashSet<RowName>,
    name_id_map: &HashMap<ColumnName, VariableID>,
    quad_terms: Option<&HashMap<(ColumnName, ColumnName), f64>>,
) -> crate::Result<(Function, Equality)> {
    let mut negate = false;

    let equality = if eq.contains(name) {
        if b != 0. {
            b = -b;
        }
        Equality::EqualToZero
    } else if le.contains(name) {
        if b != 0. {
            b = -b;
        }
        Equality::LessThanOrEqualToZero
    } else if ge.contains(name) {
        // must multiply all terms by -1
        negate = true;
        Equality::LessThanOrEqualToZero
    } else {
        // unsure what to do -- just gonna assume equality
        if b != 0. {
            b = -b;
        }
        Equality::EqualToZero
    };

    let function = if row.is_empty() && quad_terms.is_none_or(|q| q.is_empty()) {
        Function::try_from(b)?
    } else if quad_terms.is_none_or(|q| q.is_empty()) {
        // Linear constraint only
        // Build linear function by adding terms
        let mut linear = crate::Linear::try_from(b)?;

        for (col_name, &coefficient) in row {
            if let Some(&id) = name_id_map.get(col_name) {
                let coeff = if negate { -coefficient } else { coefficient };
                match Coefficient::try_from(coeff) {
                    Ok(coef) => {
                        linear.add_term(crate::LinearMonomial::Variable(id), coef)?;
                    }
                    Err(crate::CoefficientError::Zero) => {} // Skip zero coefficients
                    Err(e) => return Err(e.into()),
                }
            }
        }

        if linear.is_zero() {
            Function::Zero
        } else {
            Function::Linear(linear)
        }
    } else {
        // Quadratic constraint - need to build a quadratic function
        let mut quadratic = Quadratic::try_from(b)?;

        // Add linear terms
        for (col_name, &coefficient) in row {
            if let Some(&id) = name_id_map.get(col_name) {
                let coeff = if negate { -coefficient } else { coefficient };
                match Coefficient::try_from(coeff) {
                    Ok(coef) => {
                        quadratic.add_term(crate::QuadraticMonomial::Linear(id), coef)?;
                    }
                    Err(crate::CoefficientError::Zero) => {} // Skip zero coefficients
                    Err(e) => return Err(e.into()),
                }
            }
        }

        // Add quadratic terms
        if let Some(quad_terms) = quad_terms {
            for ((var1_name, var2_name), &coefficient) in quad_terms {
                if let (Some(&id1), Some(&id2)) =
                    (name_id_map.get(var1_name), name_id_map.get(var2_name))
                {
                    let coeff = if negate { -coefficient } else { coefficient };
                    match Coefficient::try_from(coeff) {
                        Ok(coef) => {
                            quadratic
                                .add_term(crate::QuadraticMonomial::new_pair(id1, id2), coef)?;
                        }
                        Err(crate::CoefficientError::Zero) => {} // Skip zero coefficients
                        Err(e) => return Err(e.into()),
                    }
                }
            }
        }

        Function::Quadratic(quadratic)
    };

    Ok((function, equality))
}

fn convert_sense(sense: ObjSense) -> Sense {
    match sense {
        ObjSense::Min => Sense::Minimize,
        ObjSense::Max => Sense::Maximize,
    }
}

fn get_dvar_kind(
    name: &ColumnName,
    integer: &HashSet<ColumnName>,
    binary: &HashSet<ColumnName>,
) -> DecisionVariableKind {
    if integer.contains(name) {
        DecisionVariableKind::Integer
    } else if binary.contains(name) {
        DecisionVariableKind::Binary
    } else {
        DecisionVariableKind::Continuous
    }
}

fn get_dvar_bound(
    var_name: &ColumnName,
    l: &HashMap<ColumnName, f64>,
    u: &HashMap<ColumnName, f64>,
) -> Bound {
    let (lower, upper) = match (l.get(var_name), u.get(var_name)) {
        (Some(&lower), None) => (lower, f64::INFINITY),
        (None, Some(&upper)) => {
            if upper <= 0.0 {
                (f64::NEG_INFINITY, upper)
            } else {
                (0.0, upper)
            }
        }
        (Some(&lower), Some(&upper)) => (lower, upper),

        (None, None) => (0.0, f64::INFINITY),
    };
    Bound::new(lower, upper).expect("Invalid bound")
}