pounce-nl 0.10.0

//! Minimal AMPL `.sol`-format writer.
//!
//! Format reference: David M. Gay, "Hooking Your Solver to AMPL"
//! (<https://ampl.com/REFS/hooking2.pdf>) §5 ("Returning Results to
//! AMPL"), cross-checked against the AMPL solver-library reference
//! implementation `write_sol_ASL` in
//! <https://github.com/ampl/asl> (`solvers/writesol.c`). We emit the
//! ASCII variant — the same one AMPL's `commands` file produces by
//! default when reading back from solvers.
//!
//! # Format
//!
//! ```text
//! <message line 1>
//! <message line 2>
//! ...                         (free text, ended by a blank line then "Options")
//!
//! Options
//! <nopts>                     (int — number of integer option-words to follow)
//! <opt0>                      (... nopts lines)
//! ...
//! <n_dual>                    (number of dual values written below)
//! <m>                         (constraint count)
//! <n_primal>                  (number of primal values written below)
//! <n>                         (variable count)
//! <lambda[0]>                 (... n_dual lines, dual values)
//! ...
//! <x[0]>                      (... n_primal lines, primal values)
//! ...
//! objno <objno> <status>      (optional — selects which objective and the solver-return code)
//! suffix <kind> <nvalues> <namelen> <tablen> <tabline>  (optional — one block per exported suffix)
//! <name>                      (the suffix name, on its own line)
//! <idx> <value>
//! ...
//! ```
//!
//! The four-integer count block is the canonical AMPL form: each
//! dimension count is paired with a "values written" partner so the
//! reader knows how many dual and primal lines to consume before
//! reaching `objno`. We always write every dual and primal, so
//! `n_dual == m` and `n_primal == n`. (Earlier pounce builds emitted
//! only the two bare counts `<m>\n<n>\n`; AMPL's own reader and
//! Pyomo's `.sol` reader both reject that short form.)
//!
//! # Scope
//!
//! Smallest writer that lets [`crate::nl_reader::NlSuffixes`] flow
//! from a pounce solve back through AMPL's reader. Specifically the
//! `pounce_sens` binary (pounce#17) writes:
//! * The nominal primal and dual blocks (so AMPL sees `x*` and `λ*`
//!   on the regular `_var.X` / `_con.dual` slots).
//! * One or more sensitivity suffixes (`sens_sol_state_<N>`) carrying
//!   the perturbed primal as a real-var suffix, matching upstream
//!   `MetadataMeasurement::SetSolution`
//!   (`ref/Ipopt/contrib/sIPOPT/src/SensMetadataMeasurement.cpp:128-150`).

use pounce_common::types::{Index, Number};
use std::fmt::Write as _;
use std::path::Path;

/// A single suffix block to write back into the `.sol` file. Mirrors
/// the `S`-segment shape of [`crate::nl_reader::NlSuffixes`] entries.
#[derive(Debug, Clone)]
pub struct SolSuffix {
    /// `name` as it appears in AMPL.
    pub name: String,
    /// Which side the suffix attaches to. Mapped to AMPL's
    /// `ASL_Sufkind_var` / `_con` / `_obj` / `_prob` (= 0/1/2/3).
    pub target: SolSuffixTarget,
    /// Real or integer-typed values. AMPL's `ASL_Sufkind_real` flag
    /// (`0x4`) on the kind byte selects this; we accept either typed
    /// payload here and tag the kind accordingly on write.
    pub values: SolSuffixValues,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SolSuffixTarget {
    Var = 0,
    Con = 1,
    Obj = 2,
    Problem = 3,
}

#[derive(Debug, Clone)]
pub enum SolSuffixValues {
    /// One entry per dimension of the target (variables / constraints /
    /// objectives). Sparse zero-trim happens on write — only non-zero
    /// entries land in the output, matching how AMPL emits suffixes.
    Int(Vec<Index>),
    Real(Vec<Number>),
    /// Problem-level scalar (target = Problem). Always emitted (no
    /// sparse trim, since there's only one slot).
    ProblemInt(Index),
    ProblemReal(Number),
}

/// Solution payload bundled for a `.sol` write.
#[derive(Debug, Clone)]
pub struct SolutionFile<'a> {
    /// Free-text banner / status line(s). Goes at the top of the file.
    pub message: &'a str,
    /// Primal variable values, length `n`.
    pub x: &'a [Number],
    /// Constraint multipliers in **pounce's internal (cyipopt)
    /// convention**, length `m`: the `lambda` of
    /// `L = f + lambda' g - z_L (x - x_L) + z_U (x - x_U)`, so
    /// `d obj / d b = -lambda`.
    ///
    /// [`format_sol`] negates these on the way out, because the AMPL
    /// `.sol` dual block carries *marginal values* (`d obj / d b`),
    /// not Lagrange multipliers. Pass the internal multipliers here
    /// and let the writer do the translation — do not pre-negate at
    /// the call site.
    ///
    /// See [Gay, "Hooking Your Solver to AMPL" §5](https://ampl.com/REFS/hooking2.pdf).
    pub mult_g: &'a [Number],
    /// AMPL solver return code. Convention: 0 = solved, 100..199 =
    /// "solved with warning", 200..299 = "infeasible", 300..399 =
    /// "unbounded", 400..499 = "limit reached", 500..599 = "failure".
    /// See [Gay §5, table on p. 23](https://ampl.com/REFS/hooking2.pdf).
    pub solve_result_num: i32,
    /// Suffix blocks to emit after the primal/dual blocks. Empty when
    /// no sensitivity / reduced-Hessian outputs are populated.
    pub suffixes: &'a [SolSuffix],
}

/// Format `payload` into AMPL `.sol` ASCII text, with a generic
/// `Options` block.
///
/// Prefer [`format_sol_with_options`] whenever the originating `.nl` is
/// available: a real ASL solver echoes *that model's* option words, and
/// this entry point has none to echo.
pub fn format_sol(payload: &SolutionFile<'_>) -> String {
    format_sol_with_options(payload, &[])
}

/// Format `payload` into AMPL `.sol` ASCII text, echoing `ampl_options`.
///
/// `ampl_options` are the model's own option words, read verbatim from
/// `.nl` header line 0 (`NlProblem::ampl_options`). Ipopt and the ASL's
/// `writesol.c` echo them back unchanged rather than interpreting them,
/// and the values are per-model: `bearing_400` (`g3 1 1 0`) yields
/// `3 / 1 / 1 / 0`, while `arki0003` and `camshape_6400` (`g3 10 1 0`)
/// yield `3 / 10 / 1 / 0`. Verified against Ipopt 3.14.20 `-AMPL` output.
///
/// Pass an empty slice when there is no originating header — problems
/// built through `NlProblem::from_expressions`, the WASM entry point —
/// and a generic `3 / 1 / 1 / 0` goes out instead. That is a valid block
/// (see the count discussion below); it simply is not this model's.
pub fn format_sol_with_options(payload: &SolutionFile<'_>, ampl_options: &[i64]) -> String {
    let mut out = String::new();

    // Header: message + a blank line + the `Options` block.
    //
    // The block is a count followed by that many integer option words.
    // We used to write a count of `0`, which AMPL and Pyomo's *legacy*
    // `.sol` reader accept — but no ASL solver emits it, and Pyomo's
    // v2 reader (`pyomo.contrib.solver.solvers.asl_sol_reader`, the one
    // behind `ipopt_v2` and the future default `ipopt`) reads the first
    // two option words unconditionally, to detect the documented ASL
    // quirk where a second word of `3` means two extra words follow the
    // `z` block. It therefore asserts `n_opts >= 2` and a count of `0`
    // aborts the parse, so a `.sol` POUNCE wrote could not be loaded
    // through the modern interface at all.
    //
    // Echo the model's own option words, which is what Ipopt and the
    // ASL's `writesol.c` do — they are AMPL's flags, passed through
    // rather than interpreted. With no header to echo, fall back to a
    // generic three-word block; `1` in the second slot deliberately does
    // not trigger the `3`-quirk test above.
    for line in payload.message.lines() {
        let _ = writeln!(out, "{line}");
    }
    out.push('\n');
    out.push_str("Options\n");
    if ampl_options.len() >= 2 {
        let _ = writeln!(out, "{}", ampl_options.len());
        for opt in ampl_options {
            let _ = writeln!(out, "{opt}");
        }
    } else {
        out.push_str("3\n1\n1\n0\n");
    }

    // Count block: the canonical AMPL four-integer form
    //   <n_dual_written> <n_con> <n_primal_written> <n_var>
    // The "written" counts tell the reader how many value lines to
    // consume; the bare counts are matched against the originating
    // `.nl`. We write every dual and primal, so the pairs collapse to
    // (m, m) and (n, n). Emitting only `m` and `n` (the two-integer
    // short form) makes AMPL's and Pyomo's `.sol` readers fail.
    let m = payload.mult_g.len();
    let n = payload.x.len();
    let _ = writeln!(out, "{m}");
    let _ = writeln!(out, "{m}");
    let _ = writeln!(out, "{n}");
    let _ = writeln!(out, "{n}");

    // Dual block, then primal block. AMPL writes doubles with at least
    // 16 significant digits to round-trip through IEEE-754; we use
    // Rust's `{:.17e}` to match.
    // AMPL's dual block is the *marginal value* `d obj / d b`, while
    // pounce carries the Lagrange multiplier of
    // `L = f + lambda' g`, for which `d obj / d b = -lambda`. Negate
    // on the way out so `.sol` consumers (AMPL, Pyomo `model.dual`,
    // and anything reading the file directly) see shadow prices with
    // the sign the rest of the ecosystem uses. See gh #271.
    for &v in payload.mult_g {
        let _ = writeln!(out, "{:.17e}", -v);
    }
    for &v in payload.x {
        let _ = writeln!(out, "{v:.17e}");
    }

    // Objective-number + solver return code. AMPL convention: every
    // .sol must end with at least an `objno <objno> <code>` line so
    // the reader can extract `solve_result_num`.
    let _ = writeln!(out, "objno 0 {}", payload.solve_result_num);

    // Suffix blocks. AMPL's reader skips empty / all-zero suffixes,
    // but it accepts them; we sparse-trim ints/reals to keep the
    // output small. Problem-level kinds always write a single entry.
    for s in payload.suffixes {
        write_suffix(&mut out, s);
    }

    out
}

fn write_suffix(out: &mut String, s: &SolSuffix) {
    let target_bits = s.target as u32 & 0x3;
    match &s.values {
        SolSuffixValues::Int(vs) => {
            let entries: Vec<(usize, Index)> = vs
                .iter()
                .enumerate()
                .filter(|&(_, &v)| v != 0)
                .map(|(i, &v)| (i, v))
                .collect();
            write_suffix_header(out, target_bits, entries.len(), &s.name);
            for (i, v) in entries {
                let _ = writeln!(out, "{i} {v}");
            }
        }
        SolSuffixValues::Real(vs) => {
            let entries: Vec<(usize, Number)> = vs
                .iter()
                .enumerate()
                .filter(|&(_, &v)| v != 0.0)
                .map(|(i, &v)| (i, v))
                .collect();
            write_suffix_header(out, target_bits | 0x4, entries.len(), &s.name);
            for (i, v) in entries {
                let _ = writeln!(out, "{i} {v:.17e}");
            }
        }
        SolSuffixValues::ProblemInt(v) => {
            write_suffix_header(out, target_bits, 1, &s.name);
            let _ = writeln!(out, "0 {v}");
        }
        SolSuffixValues::ProblemReal(v) => {
            write_suffix_header(out, target_bits | 0x4, 1, &s.name);
            let _ = writeln!(out, "0 {v:.17e}");
        }
    }
}

/// Emit the canonical AMPL `.sol` suffix header: five integers
/// `suffix <kind> <nvalues> <namelen> <tablen> <tabline>` followed by
/// the suffix name on its own line. `namelen` is `strlen(name)+1` (the
/// value ASL's `writesol.c` writes); `tablen`/`tabline` are 0 — pounce
/// never emits a suffix value-table. AMPL's and Pyomo's `.sol` readers
/// both require this five-integer form and read the name from the next
/// line; the older three-token `suffix <kind> <nvalues> <name>` shape
/// is rejected.
fn write_suffix_header(out: &mut String, kind: u32, nvalues: usize, name: &str) {
    let namelen = name.len() + 1;
    let _ = writeln!(out, "suffix {kind} {nvalues} {namelen} 0 0");
    let _ = writeln!(out, "{name}");
}

/// Convenience: write `payload` to `path` (truncating any existing
/// file). Returns the bytes written on success.
pub fn write_sol_file(path: &Path, payload: &SolutionFile<'_>) -> std::io::Result<usize> {
    write_sol_file_with_options(path, payload, &[])
}

/// As [`write_sol_file`], echoing the model's own `ampl_options` (see
/// [`format_sol_with_options`]).
pub fn write_sol_file_with_options(
    path: &Path,
    payload: &SolutionFile<'_>,
    ampl_options: &[i64],
) -> std::io::Result<usize> {
    let s = format_sol_with_options(payload, ampl_options);
    std::fs::write(path, &s)?;
    Ok(s.len())
}

#[cfg(test)]
mod tests {
    use super::*;

    fn options_block(text: &str) -> Vec<String> {
        let mut lines = text.lines().skip_while(|l| *l != "Options");
        lines.next();
        let n: usize = lines.next().unwrap().parse().unwrap();
        std::iter::once(n.to_string())
            .chain(lines.take(n).map(str::to_string))
            .collect()
    }

    /// A solver echoes the model's own AMPL option words rather than
    /// inventing them. `arki0003` and `camshape_6400` carry `g3 10 1 0`,
    /// and Ipopt 3.14.20 writes `3 / 10 / 1 / 0` for both — verified
    /// against real `-AMPL` output, not just the spec.
    #[test]
    fn options_block_echoes_the_models_own_words() {
        let payload = SolutionFile {
            message: "POUNCE: SolveSucceeded",
            x: &[1.0],
            mult_g: &[],
            solve_result_num: 0,
            suffixes: &[],
        };
        let text = format_sol_with_options(&payload, &[10, 1, 0]);
        assert_eq!(options_block(&text), ["3", "10", "1", "0"]);
    }

    /// With no originating header, a generic block goes out. It must
    /// still satisfy Pyomo's v2 reader: at least two words, and a second
    /// word that is not the `3` that signals two extra trailing values.
    #[test]
    fn options_block_falls_back_when_there_is_no_header() {
        let payload = SolutionFile {
            message: "POUNCE: SolveSucceeded",
            x: &[1.0],
            mult_g: &[],
            solve_result_num: 0,
            suffixes: &[],
        };
        for opts in [vec![], vec![7]] {
            let text = format_sol_with_options(&payload, &opts);
            let block = options_block(&text);
            assert_eq!(block, ["3", "1", "1", "0"], "opts={opts:?}");
            assert_ne!(block[2], "3", "must not trigger the vbtol quirk");
        }
    }

    #[test]
    fn writes_basic_primal_dual_block() {
        let payload = SolutionFile {
            message: "POUNCE: SolveSucceeded",
            x: &[1.0, 2.5, -0.5],
            mult_g: &[0.1, -0.2],
            solve_result_num: 0,
            suffixes: &[],
        };
        let s = format_sol(&payload);
        // Header banner present.
        assert!(s.starts_with("POUNCE: SolveSucceeded\n"));
        // The ASL option block, as `writesol.c` and Ipopt emit it. A
        // count of `0` here trips `assert n_opts >= 2` in Pyomo's v2
        // `.sol` reader, which is what `ipopt_v2` (and the future
        // default `ipopt`) uses — see `format_sol`.
        assert!(s.contains("\nOptions\n3\n1\n1\n0\n"), "option block:\n{s}");
        // Four-integer count block: n_dual=2, m=2, n_primal=3, n=3.
        assert!(s.contains("\n2\n2\n3\n3\n"), "counts missing:\n{s}");
        // First dual line: 0.1 in exponent form.
        assert!(
            s.contains("1.00000000000000006e-1\n") || s.contains("1.0e-1\n"),
            "lambda not present:\n{s}",
        );
        // objno tail present.
        assert!(s.trim_end().ends_with("objno 0 0"));
    }

    #[test]
    fn writes_real_var_suffix_sparse_trimming_zeros() {
        let payload = SolutionFile {
            message: "POUNCE-SENS",
            x: &[0.0, 0.0],
            mult_g: &[],
            solve_result_num: 0,
            suffixes: &[SolSuffix {
                name: "sens_sol_state_1".into(),
                target: SolSuffixTarget::Var,
                // Dense (0, 5.0, 0, 3.5); only indices 1 and 3 should
                // appear.
                values: SolSuffixValues::Real(vec![0.0, 5.0, 0.0, 3.5]),
            }],
        };
        let s = format_sol(&payload);
        // Canonical header: kind = 0|0x4 = 4 (real var), 2 values,
        // namelen = 17 ("sens_sol_state_1" + NUL), no table; name on
        // the following line.
        assert!(
            s.contains("\nsuffix 4 2 17 0 0\nsens_sol_state_1\n"),
            "missing suffix header:\n{s}",
        );
        // entries present with correct indices.
        assert!(s.contains("\n1 5.0"), "missing entry idx 1:\n{s}");
        assert!(s.contains("\n3 3.5"), "missing entry idx 3:\n{s}");
        // index 0 / 2 are zero — must not appear in the suffix block.
        // (The single-digit `0` could appear elsewhere, so we anchor.)
        assert!(!s.contains("\n0 0.0"), "zero entry was not trimmed:\n{s}",);
    }

    #[test]
    fn writes_int_constraint_suffix() {
        let payload = SolutionFile {
            message: "msg",
            x: &[],
            mult_g: &[],
            solve_result_num: 0,
            suffixes: &[SolSuffix {
                name: "sens_init_constr".into(),
                target: SolSuffixTarget::Con,
                values: SolSuffixValues::Int(vec![0, 1, 2, 0]),
            }],
        };
        let s = format_sol(&payload);
        // kind = 1 (con, integer), 2 values, namelen = 17.
        assert!(s.contains("\nsuffix 1 2 17 0 0\nsens_init_constr\n"), "{s}");
        assert!(s.contains("\n1 1\n"));
        assert!(s.contains("\n2 2\n"));
    }

    #[test]
    fn writes_problem_real_suffix() {
        let payload = SolutionFile {
            message: "msg",
            x: &[],
            mult_g: &[],
            solve_result_num: 0,
            suffixes: &[SolSuffix {
                name: "wall_time".into(),
                target: SolSuffixTarget::Problem,
                values: SolSuffixValues::ProblemReal(0.0123),
            }],
        };
        let s = format_sol(&payload);
        // kind = 3 | 0x4 = 7 (problem-level, real), namelen = 10.
        assert!(s.contains("\nsuffix 7 1 10 0 0\nwall_time\n"), "{s}");
        // Single entry at idx 0.
        assert!(s.contains("0 1.23"));
    }

    #[test]
    fn round_trip_through_nl_reader_suffix_parser() {
        // Build a .sol with an integer var-suffix, then feed the
        // suffix block to the .nl-style parser to confirm shape /
        // index conventions agree. We don't reuse parse_nl_text here
        // because the .sol prefix differs from .nl; instead we just
        // string-search the emitted suffix header against the
        // {kind, name, count} contract.
        let payload = SolutionFile {
            message: "m",
            x: &[],
            mult_g: &[],
            solve_result_num: 0,
            suffixes: &[SolSuffix {
                name: "foo".into(),
                target: SolSuffixTarget::Var,
                values: SolSuffixValues::Int(vec![1, 0, 3]),
            }],
        };
        let s = format_sol(&payload);
        // kind = 0 (var int), 2 values, namelen = 4.
        assert!(s.contains("\nsuffix 0 2 4 0 0\nfoo\n"), "{s}");
    }
}