gam_terms/inference/

formula_dsl.rs

use std::collections::BTreeMap;

use pest::Parser;
use pest::iterators::Pair;
use pest_derive::Parser;

use crate::smooth::BoundedCoefficientPriorSpec;
use gam_problem::types::{
    InverseLink, LikelihoodSpec, LinkComponent, LinkFunction, StandardLink, WigglePenaltyConfig,
};

#[derive(Parser)]
#[grammar_inline = r#"
WHITESPACE = _{ " " | "\t" | NEWLINE }

top_function_call = { SOI ~ function_call ~ EOI }
top_expr = { SOI ~ expr ~ EOI }
formula = { SOI ~ expr ~ "~" ~ rhs ~ EOI }
rhs = { term ~ ("+" ~ term)* }
term = { expr }

expr = { sum }
sum = { product ~ (add_op ~ product)* }
add_op = { "+" | "-" }
product = { interact ~ (mul_op ~ interact)* }
mul_op = { "*" | "/" }
interact = { power ~ (interact_op ~ power)* }
interact_op = { ":" }
power = { unary ~ (pow_op ~ unary)* }
pow_op = { "^" }
unary = { unary_op* ~ primary }
unary_op = _{ "+" | "-" }

primary = { function_call | list_lit | tuple_lit | ident | number | string_lit | "(" ~ expr ~ ")" }
list_lit = @{ "[" ~ (!"]" ~ ANY)* ~ "]" }
tuple_lit = @{ "(" ~ (!("," | ")") ~ ANY)+ ~ "," ~ (!")" ~ ANY)* ~ ")" }
function_call = { ident ~ "(" ~ arg_list? ~ ")" }
arg_list = { arg ~ ("," ~ arg)* }
arg = { named_arg | expr }
named_arg = { ident ~ "=" ~ expr }

ident = @{ ident_start ~ ident_continue* }
ident_start = _{ ASCII_ALPHA | "_" }
ident_continue = _{ ASCII_ALPHANUMERIC | "_" | "." }

number = @{
    "-"?
    ~ (ASCII_DIGIT+ ~ ("." ~ ASCII_DIGIT*)? | "." ~ ASCII_DIGIT+)
    ~ (("e" | "E") ~ ("+" | "-")? ~ ASCII_DIGIT+)?
}

string_lit = @{ "\"" ~ (!"\"" ~ ANY)* ~ "\"" | "'" ~ (!"'" ~ ANY)* ~ "'" }
"#]
struct FormulaParser;

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct FormulaDslParse {
    pub response_expr: String,
    pub rhs_terms: Vec<String>,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum CallArgSpec {
    Positional(String),
    Named { key: String, value: String },
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct FunctionCallSpec {
    pub name: String,
    pub args: Vec<CallArgSpec>,
}

/// Typed error surface for the formula DSL parser.
///
/// Every variant carries a free-form `reason: String` payload; `Display`
/// emits exactly that payload, so converting a `FormulaDslError` into
/// `String` (via the `From` impl below) is byte-equivalent to the pre-
/// refactor `Err(format!(...))` / `Err("...".to_string())` strings that
/// the same call sites produced. Public entry points keep their existing
/// `Result<_, String>` signatures — CLI input handling stays unchanged —
/// and typed errors flow across the boundary via `From<FormulaDslError>
/// for String`.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum FormulaDslError {
    /// Pest grammar failure, unbalanced delimiters, empty terms, or
    /// missing required parse fragments — i.e. the formula text is
    /// not a well-formed DSL string.
    ParseError { reason: String },
    /// A referenced symbol (link name, blended-link component, term
    /// function name, top-level RHS identifier) is not part of the
    /// supported vocabulary.
    UnknownIdentifier { reason: String },
    /// A named option's value is unparseable, out of range, or not a
    /// finite number / valid integer.
    InvalidArgument { reason: String },
    /// A combination of terms or options is disallowed (duplicate
    /// terms, multiple linkwiggle/link/survmodel, mutually exclusive
    /// option groups in bounded(), wiggle-incompatible links, etc.).
    IncompatibleTerm { reason: String },
    /// A required configuration option is missing or empty (e.g.
    /// `link()` without `type=`, `survmodel()` with no options,
    /// `bounded()` without a required argument).
    MalformedConfig { reason: String },
}

gam_linalg::impl_reason_error_boilerplate! {
    FormulaDslError {
        ParseError,
        UnknownIdentifier,
        InvalidArgument,
        IncompatibleTerm,
        MalformedConfig,
    }
}

/// Inbound conversion from `String` is used by `?` cascades inside `parse_formula`
/// and friends so that internal parser helpers still returning `Result<_, String>`
/// can flow through without each call site needing an explicit `.map_err(...)`.
/// We route into `ParseError` because by construction every internal helper that
/// still produces a raw `String` is itself a parse/term-resolution stage.
impl From<String> for FormulaDslError {
    fn from(reason: String) -> Self {
        FormulaDslError::ParseError { reason }
    }
}

pub fn parse_formula_dsl(formula: &str) -> Result<FormulaDslParse, String> {
    validate_balanced_delimiters(formula, "invalid formula syntax")?;
    let mut parsed =
        FormulaParser::parse(Rule::formula, formula).map_err(|e| FormulaDslError::ParseError {
            reason: format!("invalid formula syntax: {e}"),
        })?;
    let formula_pair = parsed.next().ok_or_else(|| FormulaDslError::ParseError {
        reason: "invalid formula syntax: empty parse".to_string(),
    })?;

    let mut response_expr: Option<String> = None;
    let mut rhs_terms: Option<Vec<String>> = None;

    for part in formula_pair.into_inner() {
        match part.as_rule() {
            Rule::expr if response_expr.is_none() => {
                response_expr = Some(part.as_str().trim().to_string());
            }
            Rule::rhs => {
                rhs_terms = Some(extract_rhs_terms(part)?);
            }
            _ => {}
        }
    }

    let response_expr = response_expr.ok_or_else(|| FormulaDslError::ParseError {
        reason: "invalid formula: missing response expression".to_string(),
    })?;
    let rhs_terms = rhs_terms.ok_or_else(|| FormulaDslError::ParseError {
        reason: "invalid formula: missing RHS terms".to_string(),
    })?;
    if rhs_terms.is_empty() {
        return Err(FormulaDslError::ParseError {
            reason: "formula has no usable terms".to_string(),
        }
        .into());
    }

    Ok(FormulaDslParse {
        response_expr,
        rhs_terms,
    })
}

fn delimiter_balance_error(prefix: &str) -> String {
    format!("{prefix}: unbalanced parentheses or quotes")
}

// Pest reports malformed delimiters as a generic parse failure. We validate the
// raw text first so callers get a stable, specific error class for unmatched
// parentheses/quotes instead of whichever grammar branch happened to fail last.
fn validate_balanced_delimiters(input: &str, prefix: &str) -> Result<(), String> {
    let mut stack = Vec::<char>::new();
    let mut in_single = false;
    let mut in_double = false;

    for ch in input.chars() {
        match ch {
            '\'' if !in_double => in_single = !in_single,
            '"' if !in_single => in_double = !in_double,
            '(' | '[' | '{' if !in_single && !in_double => stack.push(ch),
            ')' | ']' | '}' if !in_single && !in_double => {
                let expected = match ch {
                    ')' => '(',
                    ']' => '[',
                    // The outer match arm guarantees ch is one of ')', ']', '}'.
                    _ => '{',
                };
                if stack.pop() != Some(expected) {
                    return Err(FormulaDslError::ParseError {
                        reason: delimiter_balance_error(prefix),
                    }
                    .into());
                }
            }
            _ => {}
        }
    }

    if in_single || in_double || !stack.is_empty() {
        return Err(FormulaDslError::ParseError {
            reason: delimiter_balance_error(prefix),
        }
        .into());
    }
    Ok(())
}

fn extract_rhs_terms(rhs: Pair<'_, Rule>) -> Result<Vec<String>, String> {
    let mut out = Vec::new();
    let mut depth = 0_i32;
    let mut in_single = false;
    let mut in_double = false;
    let mut start = 0_usize;
    // Last non-whitespace character seen at depth 0 outside quotes. A top-level
    // `+` only separates terms when it follows a completed operand; when it
    // follows a binary operator (`:`, `*`, `/`, `^`) or another sign — or opens
    // the RHS — it is a UNARY sign, not a separator. Splitting on it there would
    // hand the grammar a truncated fragment like `x:` (from `x:+z`) and surface
    // a confusing "invalid term syntax in `x:`" instead of the dedicated unary
    // diagnostic the grammar raises when the whole `x:+z` term reaches it. This
    // never suppresses a split in a valid formula: no valid Wilkinson-Rogers RHS
    // places `+` immediately after an operator. `-` is already never treated as
    // a separator here (it is rejected downstream as a binary term operator), so
    // the unary `-` path already flows through intact — this restores the same
    // intact flow for unary `+`.
    let mut last_significant: Option<char> = None;
    let text = rhs.as_str();
    let bytes = text.as_bytes();
    for (idx, &b) in bytes.iter().enumerate() {
        let ch = b as char;
        match ch {
            '\'' if !in_double => in_single = !in_single,
            '"' if !in_single => in_double = !in_double,
            '(' | '[' | '{' if !in_single && !in_double => depth += 1,
            ')' | ']' | '}' if !in_single && !in_double && depth > 0 => depth -= 1,
            '+' if !in_single
                && !in_double
                && depth == 0
                && !matches!(
                    last_significant,
                    None | Some(':' | '*' | '/' | '^' | '+' | '-')
                ) =>
            {
                let term = text[start..idx].trim();
                if term.is_empty() {
                    return Err(FormulaDslError::ParseError {
                        reason: "formula RHS contains an empty term".to_string(),
                    }
                    .into());
                }
                out.push(term.to_string());
                start = idx + 1;
            }
            _ => {}
        }
        if !ch.is_ascii_whitespace() {
            last_significant = Some(ch);
        }
    }
    if in_single || in_double || depth != 0 {
        return Err(FormulaDslError::ParseError {
            reason: "formula RHS has unbalanced quotes or parentheses".to_string(),
        }
        .into());
    }
    let tail = text[start..].trim();
    if tail.is_empty() {
        return Err(FormulaDslError::ParseError {
            reason: "formula RHS contains an empty term".to_string(),
        }
        .into());
    }
    out.push(tail.to_string());
    Ok(out)
}

/// Wilkinson-Rogers operator-family expansion.
///
/// A raw RHS term (already split on top-level `+`) may use the documented
/// formula operators `:`, `*`, `/`, `^`, and parenthesization. This function
/// expands the term into a normalized list of `WrAtomList`s, where each
/// `WrAtomList` is one resulting model term — either a single atom (linear
/// main effect / function-call term) or multiple atoms (interaction).
///
/// Wilkinson-Rogers semantics implemented here:
/// * `a` produces `{a}` — one main effect.
/// * `a:b` produces `{a:b}` — one interaction.
/// * `a*b` produces `{a, b, a:b}` — crossing (expanded).
/// * `a/b` produces `{a, a:b}` — nesting.
/// * `(a + b + ... )^n` produces every non-empty subset of `{a, b, ...}` of
///   size at most `n`, each subset being one interaction.
/// * `+` unions two term sets; `-` is rejected.
///
/// Atoms inside the AST may be bare identifiers or function calls. Function
/// calls are opaque — they pass through as-is and cannot participate in an
/// `:` interaction with other atoms (smooths/factors require dedicated
/// constructors like `te()`).
type WrAtomList = Vec<String>;

fn expand_wr_term(raw: &str) -> Result<Vec<WrAtomList>, String> {
    let mut parsed = FormulaParser::parse(Rule::top_expr, raw).map_err(|e| {
        FormulaDslError::ParseError {
            reason: format!("invalid term syntax in `{raw}`: {e}"),
        }
        .to_string()
    })?;
    let top = parsed.next().ok_or_else(|| {
        FormulaDslError::ParseError {
            reason: format!("invalid term syntax in `{raw}`: empty parse"),
        }
        .to_string()
    })?;
    let expr = top
        .into_inner()
        .find(|p| p.as_rule() == Rule::expr)
        .ok_or_else(|| {
            FormulaDslError::ParseError {
                reason: format!("invalid term syntax in `{raw}`: missing expr"),
            }
            .to_string()
        })?;
    let interactions = expand_expr(expr, raw)?;
    let normalized: Vec<WrAtomList> = interactions
        .into_iter()
        .map(normalize_interaction)
        .collect();
    // Deduplicate while preserving order: a*b yields {a, b, a:b}; a*b + a
    // would otherwise list `a` twice.
    let mut seen = std::collections::BTreeSet::<Vec<String>>::new();
    let mut out = Vec::<WrAtomList>::new();
    for term in normalized {
        let key = term.clone();
        if seen.insert(key) {
            out.push(term);
        }
    }
    Ok(out)
}

fn normalize_interaction(mut atoms: WrAtomList) -> WrAtomList {
    atoms.sort();
    atoms.dedup();
    atoms
}

fn expand_expr(pair: Pair<'_, Rule>, raw: &str) -> Result<Vec<WrAtomList>, String> {
    match pair.as_rule() {
        Rule::expr => {
            let inner = pair.into_inner().next().ok_or_else(|| {
                FormulaDslError::ParseError {
                    reason: format!("invalid term syntax in `{raw}`: empty expr"),
                }
                .to_string()
            })?;
            expand_expr(inner, raw)
        }
        Rule::sum => {
            let mut iter = pair.into_inner();
            let first = iter.next().ok_or_else(|| {
                FormulaDslError::ParseError {
                    reason: format!("invalid term syntax in `{raw}`: empty sum"),
                }
                .to_string()
            })?;
            let mut acc = expand_expr(first, raw)?;
            while let Some(op) = iter.next() {
                if op.as_rule() != Rule::add_op {
                    return Err(FormulaDslError::ParseError {
                        reason: format!(
                            "invalid term syntax in `{raw}`: expected add operator, got `{:?}`",
                            op.as_rule()
                        ),
                    }
                    .into());
                }
                let op_str = op.as_str().trim();
                let operand = iter.next().ok_or_else(|| {
                    FormulaDslError::ParseError {
                        reason: format!("invalid term syntax in `{raw}`: dangling `{op_str}`"),
                    }
                    .to_string()
                })?;
                if op_str == "-" {
                    return Err(FormulaDslError::IncompatibleTerm {
                        reason: format!(
                            "binary `-` is not supported inside a formula term in `{raw}` \
                             (use multiple `+` terms or drop the unwanted predictor explicitly)"
                        ),
                    }
                    .into());
                }
                let mut rhs = expand_expr(operand, raw)?;
                acc.append(&mut rhs);
            }
            Ok(acc)
        }
        Rule::product => {
            let mut iter = pair.into_inner();
            let first = iter.next().ok_or_else(|| {
                FormulaDslError::ParseError {
                    reason: format!("invalid term syntax in `{raw}`: empty product"),
                }
                .to_string()
            })?;
            let mut acc = expand_expr(first, raw)?;
            while let Some(op) = iter.next() {
                if op.as_rule() != Rule::mul_op {
                    return Err(FormulaDslError::ParseError {
                        reason: format!(
                            "invalid term syntax in `{raw}`: expected `*` or `/`, got `{:?}`",
                            op.as_rule()
                        ),
                    }
                    .into());
                }
                let op_str = op.as_str().trim();
                let operand = iter.next().ok_or_else(|| {
                    FormulaDslError::ParseError {
                        reason: format!("invalid term syntax in `{raw}`: dangling `{op_str}`"),
                    }
                    .to_string()
                })?;
                let rhs = expand_expr(operand, raw)?;
                acc = match op_str {
                    "*" => wr_cross(acc, rhs),
                    "/" => wr_nest(acc, rhs),
                    other => {
                        return Err(FormulaDslError::ParseError {
                            reason: format!(
                                "invalid term syntax in `{raw}`: unrecognized mul operator `{other}`"
                            ),
                        }
                        .into());
                    }
                };
            }
            Ok(acc)
        }
        Rule::interact => {
            let mut iter = pair.into_inner();
            let first = iter.next().ok_or_else(|| {
                FormulaDslError::ParseError {
                    reason: format!("invalid term syntax in `{raw}`: empty interact"),
                }
                .to_string()
            })?;
            let mut acc = expand_expr(first, raw)?;
            while let Some(op) = iter.next() {
                if op.as_rule() != Rule::interact_op {
                    return Err(FormulaDslError::ParseError {
                        reason: format!(
                            "invalid term syntax in `{raw}`: expected `:`, got `{:?}`",
                            op.as_rule()
                        ),
                    }
                    .into());
                }
                let operand = iter.next().ok_or_else(|| {
                    FormulaDslError::ParseError {
                        reason: format!("invalid term syntax in `{raw}`: dangling `:`"),
                    }
                    .to_string()
                })?;
                let rhs = expand_expr(operand, raw)?;
                acc = wr_interact(acc, rhs, raw)?;
            }
            Ok(acc)
        }
        Rule::power => {
            let mut iter = pair.into_inner();
            let first = iter.next().ok_or_else(|| {
                FormulaDslError::ParseError {
                    reason: format!("invalid term syntax in `{raw}`: empty power"),
                }
                .to_string()
            })?;
            let base = expand_expr(first, raw)?;
            let Some(op) = iter.next() else {
                return Ok(base);
            };
            if op.as_rule() != Rule::pow_op {
                return Err(FormulaDslError::ParseError {
                    reason: format!(
                        "invalid term syntax in `{raw}`: expected `^`, got `{:?}`",
                        op.as_rule()
                    ),
                }
                .into());
            }
            let exponent_pair = iter.next().ok_or_else(|| {
                FormulaDslError::ParseError {
                    reason: format!("invalid term syntax in `{raw}`: dangling `^`"),
                }
                .to_string()
            })?;
            let exp_text = exponent_pair.as_str().trim();
            let n: usize = exp_text.parse().map_err(|_| {
                FormulaDslError::ParseError {
                    reason: format!(
                        "invalid term syntax in `{raw}`: `^` exponent must be a positive integer, got `{exp_text}`"
                    ),
                }
                .to_string()
            })?;
            if n == 0 {
                return Err(FormulaDslError::ParseError {
                    reason: format!(
                        "invalid term syntax in `{raw}`: `^0` is not a meaningful formula expansion"
                    ),
                }
                .into());
            }
            if let Some(extra_op) = iter.next() {
                let extra = extra_op.as_str().trim();
                return Err(FormulaDslError::ParseError {
                    reason: format!(
                        "invalid term syntax in `{raw}`: chained `^` operators are not supported; \
                         use one positive integer exponent, got another `{extra}`"
                    ),
                }
                .into());
            }
            Ok(wr_power(base, n))
        }
        Rule::unary => {
            let unary_start = pair.as_span().start();
            for inner in pair.into_inner() {
                if inner.as_rule() == Rule::primary {
                    let prefix_len = inner.as_span().start().saturating_sub(unary_start);
                    if prefix_len > 0 {
                        let prefix = &raw[unary_start..inner.as_span().start()];
                        if prefix.chars().any(|ch| matches!(ch, '+' | '-')) {
                            return Err(FormulaDslError::IncompatibleTerm {
                                reason: format!(
                                    "unary `+`/`-` is not supported inside a formula term in `{raw}`"
                                ),
                            }
                            .into());
                        }
                    }
                    return expand_expr(inner, raw);
                }
            }
            Err(FormulaDslError::ParseError {
                reason: format!("invalid term syntax in `{raw}`: empty unary"),
            }
            .into())
        }
        Rule::primary => {
            let span = pair.as_str().trim().to_string();
            let inner = pair.into_inner().next();
            match inner {
                Some(child) if child.as_rule() == Rule::expr => expand_expr(child, raw),
                Some(child) => Ok(vec![vec![child.as_str().trim().to_string()]]),
                None => Ok(vec![vec![span]]),
            }
        }
        _ => Err(FormulaDslError::ParseError {
            reason: format!(
                "invalid term syntax in `{raw}`: unexpected node `{:?}`",
                pair.as_rule()
            ),
        }
        .into()),
    }
}

fn wr_cross(left: Vec<WrAtomList>, right: Vec<WrAtomList>) -> Vec<WrAtomList> {
    // a*b = a + b + a:b
    let mut out = Vec::with_capacity(left.len() + right.len() + left.len() * right.len());
    out.extend(left.iter().cloned());
    out.extend(right.iter().cloned());
    for l in &left {
        for r in &right {
            let mut merged: WrAtomList = l.iter().cloned().chain(r.iter().cloned()).collect();
            merged.sort();
            merged.dedup();
            out.push(merged);
        }
    }
    out
}

fn wr_nest(left: Vec<WrAtomList>, right: Vec<WrAtomList>) -> Vec<WrAtomList> {
    // a/b = a + a:b
    let mut out = Vec::with_capacity(left.len() + left.len() * right.len());
    out.extend(left.iter().cloned());
    for l in &left {
        for r in &right {
            let mut merged: WrAtomList = l.iter().cloned().chain(r.iter().cloned()).collect();
            merged.sort();
            merged.dedup();
            out.push(merged);
        }
    }
    out
}

fn wr_interact(
    left: Vec<WrAtomList>,
    right: Vec<WrAtomList>,
    raw: &str,
) -> Result<Vec<WrAtomList>, String> {
    // a:b — Cartesian merge of every (l, r) pair.
    let mut out = Vec::with_capacity(left.len() * right.len());
    for l in &left {
        for r in &right {
            let combined: WrAtomList = l.iter().cloned().chain(r.iter().cloned()).collect();
            // Reject interactions whose atoms contain function-call syntax.
            // These would build design columns that are not simple products
            // (factors, smooths) and need a dedicated constructor.
            for atom in &combined {
                if atom.contains('(') {
                    return Err(FormulaDslError::IncompatibleTerm {
                        reason: format!(
                            "interaction operator `:` with function-call atom is not supported in `{raw}`. \
                             Use te(...) for smooth interactions or group()/factor() with a separate \
                             interaction strategy for categorical effects."
                        ),
                    }
                    .into());
                }
            }
            let mut merged = combined;
            merged.sort();
            merged.dedup();
            out.push(merged);
        }
    }
    Ok(out)
}

fn wr_power(base: Vec<WrAtomList>, n: usize) -> Vec<WrAtomList> {
    // (e)^n produces every non-empty subset of `base` of size ≤ n,
    // each subset being one interaction. The standard WR semantics
    // treat the base as a *set* of atoms — interactions inside `base`
    // are kept as-is and not re-crossed.
    if base.is_empty() {
        return Vec::new();
    }
    let m = base.len();
    let mut out = Vec::<WrAtomList>::new();
    let max_size = n.min(m);
    for size in 1..=max_size {
        // Enumerate combinations of `size` elements out of m.
        let mut indices: Vec<usize> = (0..size).collect();
        loop {
            let mut merged = WrAtomList::new();
            for &i in &indices {
                merged.extend(base[i].iter().cloned());
            }
            merged.sort();
            merged.dedup();
            out.push(merged);
            // Next combination
            let mut k = size;
            while k > 0 {
                k -= 1;
                if indices[k] != k + m - size {
                    indices[k] += 1;
                    for j in (k + 1)..size {
                        indices[j] = indices[j - 1] + 1;
                    }
                    break;
                }
                if k == 0 {
                    k = usize::MAX;
                    break;
                }
            }
            if k == usize::MAX {
                break;
            }
        }
    }
    out
}

fn is_exact_ident(raw: &str) -> bool {
    let mut chars = raw.chars();
    let Some(first) = chars.next() else {
        return false;
    };
    if !first.is_ascii_alphabetic() && first != '_' {
        return false;
    }
    chars.all(|ch| ch.is_ascii_alphanumeric() || ch == '_' || ch == '.')
}

pub fn parse_function_call(input: &str) -> Result<FunctionCallSpec, String> {
    validate_balanced_delimiters(input, "invalid function call syntax")?;
    let mut parsed = FormulaParser::parse(Rule::top_function_call, input).map_err(|e| {
        FormulaDslError::ParseError {
            reason: format!("invalid function call syntax: {e}"),
        }
    })?;
    let top = parsed.next().ok_or_else(|| FormulaDslError::ParseError {
        reason: "invalid function call syntax: empty parse".to_string(),
    })?;
    let call = top
        .into_inner()
        .find(|p| p.as_rule() == Rule::function_call)
        .ok_or_else(|| FormulaDslError::ParseError {
            reason: "invalid function call syntax: missing call".to_string(),
        })?;
    parse_call_pair(call)
}

fn parse_call_pair(call: Pair<'_, Rule>) -> Result<FunctionCallSpec, String> {
    let mut name: Option<String> = None;
    let mut args = Vec::<CallArgSpec>::new();
    for part in call.into_inner() {
        match part.as_rule() {
            Rule::ident => {
                if name.is_none() {
                    name = Some(part.as_str().trim().to_string());
                }
            }
            Rule::arg_list => {
                for a in part.into_inner() {
                    if a.as_rule() != Rule::arg {
                        continue;
                    }
                    let mut a_inner = a.into_inner();
                    let Some(first) = a_inner.next() else {
                        continue;
                    };
                    match first.as_rule() {
                        Rule::named_arg => {
                            let mut ni = first.into_inner();
                            let key = ni
                                .next()
                                .ok_or_else(|| FormulaDslError::ParseError {
                                    reason: "invalid named argument key".to_string(),
                                })?
                                .as_str()
                                .trim()
                                .to_ascii_lowercase();
                            let value = ni
                                .next()
                                .ok_or_else(|| FormulaDslError::ParseError {
                                    reason: "invalid named argument value".to_string(),
                                })?
                                .as_str()
                                .trim()
                                .to_string();
                            args.push(CallArgSpec::Named { key, value });
                        }
                        Rule::expr => {
                            args.push(CallArgSpec::Positional(first.as_str().trim().to_string()));
                        }
                        _ => {}
                    }
                }
            }
            _ => {}
        }
    }
    let name = name.ok_or_else(|| FormulaDslError::ParseError {
        reason: "invalid function call: missing name".to_string(),
    })?;
    Ok(FunctionCallSpec { name, args })
}

#[cfg(test)]
mod tests {
    use super::{
        CallArgSpec, ParsedTerm, parse_formula, parse_formula_dsl, parse_function_call,
        parse_linkwiggle_formulaspec, parsed_term_column_names, parsed_terms_reference_column,
        validate_marginal_slope_z_column_exclusion,
    };
    use std::collections::{BTreeMap, BTreeSet};

    #[test]
    fn parsed_term_column_names_includes_by_smooth_grouping_variable() {
        // The model's input contract must include a smooth's `by=` column, not
        // just its positional variable. `s(x, by=g)` consumes both `x` and `g`
        // (`term_builder` reads `options["by"]`); dropping `g` from the
        // required/consumable set would either omit a genuine predictor at fit
        // (CLI projected load) or — worse — make predict silently project the
        // `g` column away (#840 regression). A bare main effect and an
        // interaction round out the variant coverage.
        let parsed =
            parse_formula("y ~ s(x, by=g) + z + a:b").expect("formula with a by= smooth parses");
        let mut cols = BTreeSet::<String>::new();
        parsed_term_column_names(&parsed.terms, &mut cols);
        for expected in ["x", "g", "z", "a", "b"] {
            assert!(
                cols.contains(expected),
                "parsed_term_column_names dropped '{expected}'; got {cols:?}"
            );
        }
        // The response is never a *term* column (it is handled separately).
        assert!(
            !cols.contains("y"),
            "response leaked into term columns: {cols:?}"
        );
    }

    #[test]
    fn linkwiggle_parser_does_not_bake_in_cubic_only_restriction() {
        // Regression for #384: the cubic-only constraint belongs to the
        // score-warp / link-deviation `DeviationRuntime`, NOT to this shared
        // parser. `parse_linkwiggle_formulaspec` also feeds `timewiggle` and
        // the location-scale survival path, whose general monotone I-spline
        // value basis honors any `degree >= 2`. So the parser must accept
        // non-cubic degrees and carry them through verbatim; the cubic gate is
        // applied downstream only where the cubic-only runtime is built. A
        // prior fix wrongly forced `degree == 3` here, breaking timewiggle and
        // location-scale callers — this pins that the parser stays general.
        for deg in [2usize, 4, 5, 10] {
            let mut options = BTreeMap::new();
            options.insert("degree".to_string(), deg.to_string());
            options.insert("internal_knots".to_string(), "3".to_string());
            let raw = format!("timewiggle(degree={deg}, internal_knots=3)");
            let spec = parse_linkwiggle_formulaspec(&options, &raw)
                .expect("non-cubic wiggle degree must parse at the shared layer");
            assert_eq!(
                spec.degree, deg,
                "parser must carry the requested degree through verbatim"
            );
        }

        // The only universal lower bound the shared parser enforces is that a
        // polynomial degree is positive.
        let mut zero = BTreeMap::new();
        zero.insert("degree".to_string(), "0".to_string());
        zero.insert("internal_knots".to_string(), "3".to_string());
        let err = parse_linkwiggle_formulaspec(&zero, "linkwiggle(degree=0, internal_knots=3)")
            .expect_err("degree=0 must be rejected");
        assert!(
            err.contains("degree >= 1"),
            "error should state the positive-degree lower bound, got: {err}"
        );
    }

    #[test]
    fn parses_nested_formula_terms() {
        let parsed =
            parse_formula_dsl("log(y) ~ x1 + s(log(x2 + 1), bs=\"tps\", k=10) + te(x3, x4)")
                .expect("parse");
        assert_eq!(parsed.response_expr, "log(y)");
        assert_eq!(parsed.rhs_terms.len(), 3);
        assert_eq!(parsed.rhs_terms[0], "x1");
        assert_eq!(parsed.rhs_terms[1], "s(log(x2 + 1), bs=\"tps\", k=10)");
        assert_eq!(parsed.rhs_terms[2], "te(x3, x4)");
    }

    #[test]
    fn parses_cyclic_formula_aliases() {
        let parsed = parse_formula("y ~ cyclic(theta, period_start=0, period_end=6.283)")
            .expect("parse cyclic formula");
        match &parsed.terms[0] {
            super::ParsedTerm::Smooth { vars, options, .. } => {
                assert_eq!(vars, &vec!["theta".to_string()]);
                assert_eq!(options.get("type").map(String::as_str), Some("cyclic"));
                assert_eq!(options.get("period_start").map(String::as_str), Some("0"));
            }
            other => panic!("expected cyclic smooth term, got {other:?}"),
        }
    }

    #[test]
    fn sphere_aliases_all_dispatch_to_intrinsic_s2_basis() {
        // Regression for #383: `s2(lat, lon)` must build the same intrinsic
        // S² (sphere) basis as `sphere()`/`sos()`/`spherical()`. Previously the
        // `s2` arm returned a Smooth without `type=sphere`, so it silently fell
        // back to a generic Euclidean 2-D smooth and diverged in the
        // spatial-kappa optimizer. All four aliases must be byte-for-byte
        // equivalent in their dispatch (vars + `type=sphere`).
        for alias in ["sphere", "sos", "spherical", "s2"] {
            let parsed = parse_formula(&format!("y ~ {alias}(lat, lon)"))
                .unwrap_or_else(|e| panic!("parse {alias}: {e}"));
            match &parsed.terms[0] {
                super::ParsedTerm::Smooth { vars, options, .. } => {
                    assert_eq!(
                        vars,
                        &vec!["lat".to_string(), "lon".to_string()],
                        "{alias} should keep (lat, lon) as its variables"
                    );
                    assert_eq!(
                        options.get("type").map(String::as_str),
                        Some("sphere"),
                        "{alias} must dispatch to the intrinsic sphere basis (type=sphere)"
                    );
                }
                other => panic!("expected sphere smooth term for {alias}, got {other:?}"),
            }
        }
    }

    #[test]
    fn parses_function_callwithnamed_and_positional_args() {
        let call = parse_function_call("s(log(x + 1), type=\"duchon\", centers=12)").expect("call");
        assert_eq!(call.name, "s");
        assert_eq!(call.args.len(), 3);
        assert_eq!(
            call.args[0],
            CallArgSpec::Positional("log(x + 1)".to_string())
        );
        assert_eq!(
            call.args[1],
            CallArgSpec::Named {
                key: "type".to_string(),
                value: "\"duchon\"".to_string()
            }
        );
    }

    #[test]
    fn parses_tensor_boundary_list_options() {
        let call = parse_function_call(
            "te(day_of_week, hour, boundary=['periodic', 'periodic'], period=[7, 24])",
        )
        .expect("call");
        assert_eq!(call.name, "te");
        assert_eq!(call.args.len(), 4);
        assert_eq!(
            call.args[2],
            CallArgSpec::Named {
                key: "boundary".to_string(),
                value: "['periodic', 'periodic']".to_string(),
            }
        );
    }

    #[test]
    fn parse_formula_dsl_reports_unbalanced_parentheses() {
        let err = parse_formula_dsl("y ~ s(x, k=10").expect_err("expected parse failure");
        assert!(err.contains("unbalanced parentheses"));
    }

    #[test]
    fn parse_function_call_reports_unbalanced_parentheses() {
        let err = parse_function_call("s(x, k=10").expect_err("expected parse failure");
        assert!(err.contains("unbalanced parentheses"));
    }

    #[test]
    fn parse_formula_accepts_tuple_smooth_options() {
        let parsed = parse_formula("z ~ te(x, y, k=(20, 20))")
            .expect("tuple-valued smooth option should parse");
        assert_eq!(parsed.terms.len(), 1);

        let dsl = parse_formula_dsl("z ~ te(x, y, k=(20, 20))")
            .expect("tuple-valued smooth option should parse in the DSL layer");
        assert_eq!(dsl.rhs_terms, vec!["te(x, y, k=(20, 20))"]);

        let call = parse_function_call("te(x, y, k=(20, 20))")
            .expect("tuple-valued smooth option should parse as a function call");
        assert_eq!(
            call.args[2],
            CallArgSpec::Named {
                key: "k".to_string(),
                value: "(20, 20)".to_string(),
            }
        );
    }

    #[test]
    fn parse_formula_rejects_unsupported_top_level_rhs_expressions() {
        // Binary `-`, unary `-`, bare parens, and `-1` intercept removal are
        // not supported as top-level RHS expressions; they must surface
        // through the bare-identifier check in `parse_term` with the same
        // diagnostic. WR operators `:`, `*`, `/`, and `^` are intentionally
        // supported by `expand_wr_term` and are exercised by
        // `parse_formula_supports_wr_slash_nesting` /
        // `parse_formula_supports_wr_star_crossing`; they must NOT appear in
        // this list.
        for formula in ["y ~ x - z", "y ~ -x", "y ~ (x)", "y ~ x - 1"] {
            let err = parse_formula(formula).expect_err("expected formula parse failure");
            assert!(err.to_string().contains("unsupported top-level RHS term"));
        }
    }

    /// Wilkinson-Rogers `a/b` (nesting) is documented and implemented by
    /// `expand_wr_term` to yield `{a, a:b}`. This test pins that contract end
    /// to end so a regression in either the grammar's `mul_op = { "*" | "/" }`
    /// branch or in `expand_expr`'s `Rule::product` handling surfaces here
    /// instead of slipping into a silent re-interpretation as a rejection.
    #[test]
    fn parse_formula_supports_wr_slash_nesting() {
        let parsed = parse_formula("y ~ x / z").expect("`/` is supported as WR nesting");
        assert_eq!(parsed.response, "y");
        assert_eq!(parsed.terms.len(), 2);
        let names: Vec<String> = parsed
            .terms
            .iter()
            .map(|t| match t {
                ParsedTerm::Linear { name, .. } => format!("Linear({name})"),
                ParsedTerm::Interaction { vars, .. } => {
                    format!("Interaction({})", vars.join(":"))
                }
                other => format!("Other({other:?})"),
            })
            .collect();
        assert_eq!(
            names,
            vec!["Linear(x)".to_string(), "Interaction(x:z)".to_string()]
        );
    }

    /// Wilkinson-Rogers `a*b` (crossing) is documented and implemented by
    /// `expand_wr_term` to yield `{a, b, a:b}`. Pin the contract so a future
    /// refactor of `mul_op` / `Rule::product` handling cannot silently change
    /// the set of model terms.
    #[test]
    fn parse_formula_supports_wr_star_crossing() {
        let parsed = parse_formula("y ~ x * z").expect("`*` is supported as WR crossing");
        assert_eq!(parsed.response, "y");
        assert_eq!(parsed.terms.len(), 3);
        let names: Vec<String> = parsed
            .terms
            .iter()
            .map(|t| match t {
                ParsedTerm::Linear { name, .. } => format!("Linear({name})"),
                ParsedTerm::Interaction { vars, .. } => {
                    format!("Interaction({})", vars.join(":"))
                }
                other => format!("Other({other:?})"),
            })
            .collect();
        assert_eq!(
            names,
            vec![
                "Linear(x)".to_string(),
                "Linear(z)".to_string(),
                "Interaction(x:z)".to_string(),
            ]
        );
    }

    #[test]
    fn parse_formula_rejects_unary_signs_inside_wr_expansion() {
        for formula in ["y ~ x:-z", "y ~ a*-b", "y ~ x/-z", "y ~ x:+z"] {
            let err = parse_formula(formula)
                .expect_err("WR expansion must not silently drop unary signs");
            let msg = err.to_string();
            assert!(
                msg.contains("unary `+`/`-` is not supported"),
                "unexpected error for {formula}: {msg}"
            );
        }
    }

    #[test]
    fn parse_formula_supports_wr_power_crossing() {
        let parsed = parse_formula("y ~ (x + z)^2").expect("`^` is supported as WR power");
        assert_eq!(parsed.response, "y");
        assert_eq!(parsed.terms.len(), 3);
        let names: Vec<String> = parsed
            .terms
            .iter()
            .map(|t| match t {
                ParsedTerm::Linear { name, .. } => format!("Linear({name})"),
                ParsedTerm::Interaction { vars, .. } => {
                    format!("Interaction({})", vars.join(":"))
                }
                other => format!("Other({other:?})"),
            })
            .collect();
        assert_eq!(
            names,
            vec![
                "Linear(x)".to_string(),
                "Linear(z)".to_string(),
                "Interaction(x:z)".to_string(),
            ]
        );
    }

    #[test]
    fn parse_formula_rejects_chained_wr_power() {
        let err = parse_formula("y ~ (x + z)^2^3")
            .expect_err("chained WR powers must not silently drop later exponents");
        let msg = err.to_string();
        assert!(
            msg.contains("chained `^` operators are not supported"),
            "error should explain that chained powers are rejected, got: {msg}"
        );
    }

    #[test]
    fn parsed_terms_reference_column_sees_the_by_smooth_variable() {
        // Regression for #807: the by= grouping variable lives in
        // options["by"], not the smooth's positional vars. The reference
        // predicate must still recognise it, both so the CLI loads the column
        // and so the marginal-slope z-column exclusion check (which reuses this
        // predicate) cannot be fooled into aliasing a reserved z onto a by=.
        let parsed = parse_formula("y ~ s(x, by=g)").expect("parse by-smooth");
        assert!(
            parsed_terms_reference_column(&parsed.terms, "g"),
            "s(x, by=g) references column g via options[\"by\"]"
        );
        assert!(parsed_terms_reference_column(&parsed.terms, "x"));
        assert!(!parsed_terms_reference_column(&parsed.terms, "absent"));
    }

    #[test]
    fn marginal_slope_z_column_validator_detects_linear_and_smooth_reuse() {
        let main = parse_formula("y ~ x + z").expect("parse main");
        let logslope = parse_formula("y ~ s(z, type=duchon, centers=6)").expect("parse logslope");

        assert!(parsed_terms_reference_column(&main.terms, "z"));
        assert!(parsed_terms_reference_column(&logslope.terms, "z"));

        let err = validate_marginal_slope_z_column_exclusion(
            &main,
            &parse_formula("y ~ 1").expect("parse clean logslope"),
            "z",
            "bernoulli marginal-slope",
            "--logslope-formula",
        )
        .expect_err("main formula should be rejected");
        assert!(err.contains("cannot also appear in the main formula"));

        let err = validate_marginal_slope_z_column_exclusion(
            &parse_formula("y ~ x").expect("parse clean main"),
            &logslope,
            "z",
            "bernoulli marginal-slope",
            "--logslope-formula",
        )
        .expect_err("logslope formula should be rejected");
        assert!(err.contains("cannot also appear in --logslope-formula"));
    }

    #[test]
    fn logslope_surface_declarations_are_additive() {
        let parsed = parse_formula("y ~ s(pc1) + logslope(z2, s(pc2)) + logslope(z3, x3)")
            .expect("parse additive logslope surfaces");
        assert_eq!(parsed.terms.len(), 1);
        assert_eq!(parsed.logslope_surfaces.len(), 2);
        assert_eq!(parsed.logslope_surfaces[0].z_column, "z2");
        assert_eq!(parsed.logslope_surfaces[0].terms.len(), 1);
        assert_eq!(parsed.logslope_surfaces[1].z_column, "z3");
        assert_eq!(parsed.logslope_surfaces[1].terms.len(), 1);
    }

    #[test]
    fn marginal_slope_z_column_validator_reserves_all_surface_z_columns() {
        let main = parse_formula("y ~ x").expect("parse main");
        let logslope = parse_formula("y ~ s(pc1) + logslope(z2, s(z3)) + logslope(z3, x)")
            .expect("parse logslope surfaces");
        let err = validate_marginal_slope_z_column_exclusion(
            &main,
            &logslope,
            "z1",
            "bernoulli marginal-slope",
            "--logslope-formula",
        )
        .expect_err("surface formula should reject another reserved z coordinate");
        assert!(err.contains("reserves z column 'z3'"));
    }

    /// Extract the single `RandomEffect` term's unseen-level policy from a
    /// one-term formula, panicking if the term is not a random-effect block.
    fn random_effect_lenient_unseen(formula: &str) -> bool {
        let parsed = parse_formula(formula).expect("parse random-effect formula");
        let re = parsed.terms.iter().find_map(|t| match t {
            ParsedTerm::RandomEffect { lenient_unseen, .. } => Some(*lenient_unseen),
            _ => None,
        });
        re.unwrap_or_else(|| panic!("{formula} did not lower to a RandomEffect term"))
    }

    #[test]
    fn factor_wrapper_is_strict_on_unseen_levels_while_group_re_are_lenient() {
        // Regression for #2137 (sibling of #2102): `factor(g)` is a FIXED
        // categorical factor (R `factor()` / patsy `C()`), so an out-of-vocabulary
        // level at predict is a schema mismatch that must raise — NOT be shrunk to
        // the centering point. `group(g)`/`re(g)`/`s(g, bs="re")` are genuine
        // random effects that tolerate a held-out group (→ population mean). The
        // parse arm once hardcoded `lenient_unseen: true` for all four wrappers,
        // so `factor(g)` silently averaged an unseen level. Pin the per-wrapper
        // policy at the parse layer, where the whole distinction now lives.
        assert!(
            !random_effect_lenient_unseen("y ~ factor(g)"),
            "factor(g) is a fixed categorical factor: strict (lenient_unseen=false) on unseen levels"
        );
        for lenient in ["y ~ group(g)", "y ~ re(g)", "y ~ s(g, bs=re)"] {
            assert!(
                random_effect_lenient_unseen(lenient),
                "{lenient} is a genuine random effect: lenient (lenient_unseen=true) on unseen levels"
            );
        }
    }
}

// ---------------------------------------------------------------------------
// Higher-level formula parsing: ParsedFormula, ParsedTerm, and friends
// ---------------------------------------------------------------------------

#[derive(Clone, Debug)]
pub struct LinkWiggleFormulaSpec {
    pub degree: usize,
    pub num_internal_knots: usize,
    pub penalty_orders: Vec<usize>,
    pub double_penalty: bool,
}

pub fn default_linkwiggle_formulaspec() -> LinkWiggleFormulaSpec {
    let cfg = WigglePenaltyConfig::cubic_triple_operator_default();
    LinkWiggleFormulaSpec {
        degree: cfg.degree,
        num_internal_knots: cfg.num_internal_knots,
        penalty_orders: cfg.penalty_orders,
        double_penalty: cfg.double_penalty,
    }
}

#[derive(Clone, Debug)]
pub struct LinkFormulaSpec {
    pub link: String,
    pub mixture_rho: Option<String>,
    pub sas_init: Option<String>,
    pub beta_logistic_init: Option<String>,
}

#[derive(Clone, Debug)]
pub struct SurvivalFormulaSpec {
    pub spec: Option<String>,
    pub survival_distribution: Option<String>,
}

#[derive(Clone, Debug)]
pub struct ParsedFormula {
    pub response: String,
    pub terms: Vec<ParsedTerm>,
    pub logslope_surfaces: Vec<LogSlopeSurfaceSpec>,
    pub linkwiggle: Option<LinkWiggleFormulaSpec>,
    pub timewiggle: Option<LinkWiggleFormulaSpec>,
    pub linkspec: Option<LinkFormulaSpec>,
    pub survivalspec: Option<SurvivalFormulaSpec>,
}

#[derive(Clone, Debug)]
pub struct LogSlopeSurfaceSpec {
    pub z_column: String,
    pub terms: Vec<ParsedTerm>,
}

pub fn marginal_slope_logslope_surfaces(
    logslope_formula: &ParsedFormula,
    default_z_column: &str,
) -> Result<Vec<LogSlopeSurfaceSpec>, String> {
    let mut surfaces = Vec::new();
    if !logslope_formula.terms.is_empty() {
        surfaces.push(LogSlopeSurfaceSpec {
            z_column: default_z_column.to_string(),
            terms: logslope_formula.terms.clone(),
        });
    }
    surfaces.extend(logslope_formula.logslope_surfaces.clone());
    if surfaces.is_empty() {
        surfaces.push(LogSlopeSurfaceSpec {
            z_column: default_z_column.to_string(),
            terms: Vec::new(),
        });
    }
    let mut seen = std::collections::BTreeSet::<String>::new();
    for surface in &surfaces {
        if !seen.insert(surface.z_column.clone()) {
            return Err(FormulaDslError::IncompatibleTerm {
                reason: format!(
                    "logslope formula declares z column '{}' more than once; each z coordinate needs exactly one log-slope surface",
                    surface.z_column
                ),
            }
            .into());
        }
    }
    Ok(surfaces)
}

#[derive(Clone, Debug)]
pub enum ParsedTerm {
    Linear {
        name: String,
        explicit: bool,
        double_penalty: bool,
        coefficient_min: Option<f64>,
        coefficient_max: Option<f64>,
    },
    BoundedLinear {
        name: String,
        min: f64,
        max: f64,
        prior: BoundedCoefficientPriorSpec,
        double_penalty: bool,
    },
    RandomEffect {
        name: String,
        /// Unseen-level policy, fixed at parse time by the wrapper the user
        /// wrote. `group(g)`/`re(g)`/`s(g, bs="re")` are genuine **random
        /// effects**: a held-out group is shrunk to the population mean, so an
        /// unseen level at predict is tolerated (`true`). `factor(g)` is a
        /// **fixed** categorical factor (R `factor()` / patsy `C()`
        /// convention): like a bare `+ g` categorical main effect, an unseen
        /// level is a schema mismatch that must raise rather than collapse onto
        /// the factor's centering point (`false`, #2137/#2102). Both wrappers
        /// share the penalized-categorical materialization; only this policy
        /// distinguishes them, so seen-level fits are identical.
        lenient_unseen: bool,
    },
    Smooth {
        label: String,
        vars: Vec<String>,
        kind: SmoothKind,
        options: BTreeMap<String, String>,
    },
    LinkWiggle {
        options: BTreeMap<String, String>,
    },
    TimeWiggle {
        options: BTreeMap<String, String>,
    },
    LinkConfig {
        options: BTreeMap<String, String>,
    },
    SurvivalConfig {
        options: BTreeMap<String, String>,
    },
    LogSlopeSurface {
        z_column: String,
        terms: Vec<ParsedTerm>,
    },
    /// Wilkinson-Rogers interaction term `a:b[:c...]`.
    ///
    /// `vars` is a sorted, deduplicated list of base column names. Each element
    /// must be a bare identifier — interactions with function-call atoms
    /// (smooths, factors, etc.) are rejected upstream because their design
    /// columns are not simple products. The design column is the elementwise
    /// product of the referenced numeric columns.
    Interaction {
        vars: Vec<String>,
        double_penalty: bool,
    },
}

/// Collect the names of every data column the parsed terms consume.
///
/// This is the canonical formula→columns walk shared by the fit-time and
/// predict-time required-column computations (the CLI and PyFFI surfaces both
/// route through it). It includes a smooth's positional `vars` *and* its `by=`
/// grouping/scaling column (`s(x, by=g)`), which `term_builder` reads from
/// `options["by"]` but which is not among the positional variables — omitting
/// it would drop a genuine predictor from the model's input contract.
pub fn parsed_term_column_names(
    terms: &[ParsedTerm],
    out: &mut std::collections::BTreeSet<String>,
) {
    for term in terms {
        match term {
            ParsedTerm::Linear { name, .. }
            | ParsedTerm::BoundedLinear { name, .. }
            | ParsedTerm::RandomEffect { name, .. } => {
                out.insert(name.clone());
            }
            ParsedTerm::Smooth { vars, options, .. } => {
                out.extend(vars.iter().cloned());
                if let Some(by) = options.get("by") {
                    out.insert(by.clone());
                }
            }
            ParsedTerm::Interaction { vars, .. } => {
                out.extend(vars.iter().cloned());
            }
            ParsedTerm::LinkWiggle { .. }
            | ParsedTerm::TimeWiggle { .. }
            | ParsedTerm::LinkConfig { .. }
            | ParsedTerm::SurvivalConfig { .. } => {}
            ParsedTerm::LogSlopeSurface { z_column, terms } => {
                out.insert(z_column.clone());
                parsed_term_column_names(terms, out);
            }
        }
    }
}

pub fn parsed_terms_reference_column(terms: &[ParsedTerm], column_name: &str) -> bool {
    terms.iter().any(|term| match term {
        ParsedTerm::Linear { name, .. }
        | ParsedTerm::BoundedLinear { name, .. }
        | ParsedTerm::RandomEffect { name, .. } => name == column_name,
        ParsedTerm::Smooth { vars, options, .. } => {
            vars.iter().any(|var| var == column_name)
                || options.get("by").is_some_and(|by| by == column_name)
        }
        ParsedTerm::Interaction { vars, .. } => vars.iter().any(|var| var == column_name),
        ParsedTerm::LinkWiggle { .. }
        | ParsedTerm::TimeWiggle { .. }
        | ParsedTerm::LinkConfig { .. }
        | ParsedTerm::SurvivalConfig { .. } => false,
        ParsedTerm::LogSlopeSurface { z_column, terms } => {
            z_column == column_name || parsed_terms_reference_column(terms, column_name)
        }
    })
}

pub fn validate_marginal_slope_z_column_exclusion(
    main_formula: &ParsedFormula,
    logslope_formula: &ParsedFormula,
    z_column: &str,
    context: &str,
    logslope_label: &str,
) -> Result<(), String> {
    let surfaces = marginal_slope_logslope_surfaces(logslope_formula, z_column)?;
    // The CLI/configured z column is reserved even when the log-slope formula
    // is intercept-only (`~ 1`) and therefore contributes no surface terms.
    // Explicit logslope(...) declarations may reserve additional z coordinates.
    let mut reserved_z_columns = std::collections::BTreeSet::<&str>::new();
    reserved_z_columns.insert(z_column);
    reserved_z_columns.extend(surfaces.iter().map(|surface| surface.z_column.as_str()));

    for reserved in &reserved_z_columns {
        if parsed_terms_reference_column(&main_formula.terms, reserved) {
            return Err(FormulaDslError::IncompatibleTerm {
                reason: format!(
                    "{context} reserves z column '{reserved}' as the auxiliary latent score; it cannot also appear in the main formula"
                ),
            }
            .into());
        }
    }
    for reserved in &reserved_z_columns {
        if parsed_terms_reference_column(&logslope_formula.terms, reserved) {
            return Err(FormulaDslError::IncompatibleTerm {
                reason: format!(
                    "{context} reserves z column '{reserved}' as the auxiliary latent score; it cannot also appear in {logslope_label}"
                ),
            }
            .into());
        }
        for surface in &surfaces {
            if parsed_terms_reference_column(&surface.terms, reserved) {
                return Err(FormulaDslError::IncompatibleTerm {
                    reason: format!(
                        "{context} reserves z column '{reserved}' as an auxiliary latent score; it cannot also appear in {logslope_label}"
                    ),
                }
                .into());
            }
        }
    }
    Ok(())
}

#[derive(Clone, Copy, Debug)]
pub enum SmoothKind {
    S,
    Te,
    /// Tensor smooth (`t2(...)`) using mgcv's separable penalty decomposition:
    /// the tensor coefficient space is split into marginal penalized/null-space
    /// tensor subspaces, with one smoothing parameter per non-null subspace.
    T2,
    /// Tensor *interaction* smooth (`ti(...)`): a tensor-product smooth whose
    /// marginal main effects are excluded, so the term captures only the pure
    /// interaction between its variables. Materializes through the same tensor
    /// path as [`SmoothKind::Te`] but with per-margin sum-to-zero
    /// identifiability (`TensorBSplineIdentifiability::MarginalSumToZero`).
    Ti,
}

#[derive(Clone, Copy, Debug)]
pub enum LinkMode {
    Strict,
    Flexible,
}

#[derive(Clone, Debug)]
pub struct LinkChoice {
    pub mode: LinkMode,
    pub link: LinkFunction,
    pub mixture_components: Option<Vec<LinkComponent>>,
}

// ---------------------------------------------------------------------------
// Link wiggle / link choice helpers
// ---------------------------------------------------------------------------

pub fn effectivelinkwiggle_formulaspec(
    formula_linkwiggle: Option<&LinkWiggleFormulaSpec>,
    link_choice: Option<&LinkChoice>,
) -> Option<LinkWiggleFormulaSpec> {
    formula_linkwiggle.cloned().or_else(|| {
        link_choice.and_then(|choice| {
            if matches!(choice.mode, LinkMode::Flexible) {
                Some(default_linkwiggle_formulaspec())
            } else {
                None
            }
        })
    })
}

pub const fn linkname_supports_joint_wiggle(link: LinkFunction) -> bool {
    !matches!(link, LinkFunction::Sas | LinkFunction::BetaLogistic)
}

pub const fn linkchoice_supports_joint_wiggle(choice: &LinkChoice) -> bool {
    match &choice.mixture_components {
        None => linkname_supports_joint_wiggle(choice.link),
        Some(_) => false,
    }
}

pub fn require_linkchoice_supports_joint_wiggle(
    choice: &LinkChoice,
    context: &str,
) -> Result<(), String> {
    if linkchoice_supports_joint_wiggle(choice) {
        Ok(())
    } else {
        Err(joint_wiggle_unsupported_link_message(context))
    }
}

pub const fn likelihood_spec_supports_joint_wiggle(likelihood: &LikelihoodSpec) -> bool {
    inverse_link_supports_joint_wiggle(&likelihood.link)
}

pub fn require_likelihood_spec_supports_joint_wiggle(
    likelihood: &LikelihoodSpec,
    context: &str,
) -> Result<(), String> {
    if likelihood_spec_supports_joint_wiggle(likelihood) {
        Ok(())
    } else {
        Err(joint_wiggle_unsupported_link_message(context))
    }
}

/// Family-agnostic capability of the joint link-wiggle machinery: which base
/// inverse links a monotone warp can be fit over AND reconstructed from at
/// predict time. Every state-less standard link qualifies — the warp fit and
/// its saved-model reconstruction evaluate the base inverse link purely through
/// the generic `inverse_link_jet_for_inverse_link` jet dispatch, which carries
/// LogLog and Cauchit exactly as it does Logit/Probit/CLogLog. LogLog/Cauchit
/// were previously omitted, so a binomial `flexible(loglog)`/`flexible(cauchit)`
/// fit that this-gate-agnostically *converged* (see
/// `binomial_inverse_link_supports_joint_wiggle`) then failed at predict when
/// `FittedModel::saved_link_wiggle` re-checked the saved link here (#2155). The
/// state-bearing links (SAS/BetaLogistic/Mixture/LatentCLogLog) carry fitted
/// warp/skew state of their own and are intentionally excluded.
pub const fn inverse_link_supports_joint_wiggle(link: &InverseLink) -> bool {
    matches!(
        link,
        InverseLink::Standard(StandardLink::Identity)
            | InverseLink::Standard(StandardLink::Log)
            | InverseLink::Standard(StandardLink::Logit)
            | InverseLink::Standard(StandardLink::Probit)
            | InverseLink::Standard(StandardLink::CLogLog)
            | InverseLink::Standard(StandardLink::LogLog)
            | InverseLink::Standard(StandardLink::Cauchit)
    )
}

pub fn require_inverse_link_supports_joint_wiggle(
    link: &InverseLink,
    context: &str,
) -> Result<(), String> {
    if inverse_link_supports_joint_wiggle(link) {
        Ok(())
    } else {
        Err(joint_wiggle_unsupported_link_message(context))
    }
}

/// Which binomial base links the joint link-wiggle (flexible-link) solver can
/// fit. All five standard binomial probability links qualify: the wiggle kernel
/// (`BinomialMeanWiggleFamily`) evaluates the base inverse link purely through
/// the generic `inverse_link_jet_for_inverse_link` dispatch, which carries full
/// jets for LogLog and Cauchit exactly as it does for Logit/Probit/CLogLog — so
/// `flexible(loglog)` / `flexible(cauchit)` fit through the same machinery
/// (#2155). Previously this gate listed only logit/probit/cloglog while the
/// permissive parse gate `linkname_supports_joint_wiggle` admitted loglog/cauchit,
/// so the config was accepted then aborted deep in the solver. The state-bearing
/// links (SAS/BetaLogistic/Mixture/LatentCLogLog) and identity/log stay out: the
/// warp is defined only over a fixed state-less base probability link.
pub const fn binomial_inverse_link_supports_joint_wiggle(link: &InverseLink) -> bool {
    matches!(
        link,
        InverseLink::Standard(StandardLink::Logit)
            | InverseLink::Standard(StandardLink::Probit)
            | InverseLink::Standard(StandardLink::CLogLog)
            | InverseLink::Standard(StandardLink::LogLog)
            | InverseLink::Standard(StandardLink::Cauchit)
    )
}

pub fn require_binomial_inverse_link_supports_joint_wiggle(
    link: &InverseLink,
    context: &str,
) -> Result<(), String> {
    if binomial_inverse_link_supports_joint_wiggle(link) {
        Ok(())
    } else {
        Err(FormulaDslError::IncompatibleTerm {
            reason: format!(
                "{context} does not support identity, log, latent-cloglog, SAS, BetaLogistic, or Mixture links; wiggle is only available for jointly fitted standard binomial probability links (logit/probit/cloglog/loglog/cauchit)"
            ),
        }
        .into())
    }
}

pub fn joint_wiggle_unsupported_link_message(context: &str) -> String {
    format!(
        "{context} does not support latent-cloglog, SAS, BetaLogistic, or Mixture links; wiggle is only available for jointly fitted standard links"
    )
}

// ---------------------------------------------------------------------------
// Option-map helpers (shared by formula parsing and term construction)
// ---------------------------------------------------------------------------

pub fn option_usize(map: &BTreeMap<String, String>, key: &str) -> Option<usize> {
    map.get(key).and_then(|v| v.parse::<usize>().ok())
}

/// Local sibling of `term_builder::validate_known_options` used by the
/// parser-side `linear / bounded / constrain / nonnegative / nonpositive`
/// branches (which build their `ParsedTerm` here and never enter
/// `term_builder::build_smooth_basis`). Without this, typos like
/// `bounded(x, min=0, max=1, foo=bar)` silently succeed because the
/// `foo` key was just never read.
fn validate_known_term_options(
    term_name: &str,
    options: &BTreeMap<String, String>,
    known: &[&str],
    raw: &str,
) -> Result<(), String> {
    let known_set: std::collections::BTreeSet<&&str> = known.iter().collect();
    for key in options.keys() {
        if !known_set.contains(&key.as_str()) {
            let known_sorted = {
                let mut v = known.to_vec();
                v.sort_unstable();
                v.join(", ")
            };
            let known_hint = if known.is_empty() {
                "no options".to_string()
            } else {
                format!("[{known_sorted}]")
            };
            return Err(FormulaDslError::InvalidArgument {
                reason: format!(
                    "{term_name}() does not accept option `{key}` (in `{raw}`); known options: {known_hint}"
                ),
            }
            .into());
        }
    }
    Ok(())
}

pub fn option_usize_any(map: &BTreeMap<String, String>, keys: &[&str]) -> Option<usize> {
    for key in keys {
        if let Some(v) = option_usize(map, key) {
            return Some(v);
        }
    }
    None
}

/// Strict integer option: returns `Ok(None)` if not present, `Ok(Some(n))` if
/// it parses as a non-negative integer, and `Err(msg)` if the user supplied a
/// value that isn't a valid usize (negative, decimal, garbage). Without this
/// the lenient `option_usize` silently drops invalid values and reverts to
/// the default — `k=-1` and `k=1.5` were both accepted as "k not specified"
/// instead of being flagged as user mistakes.
pub fn option_usize_strict(
    map: &BTreeMap<String, String>,
    key: &str,
) -> Result<Option<usize>, String> {
    match map.get(key) {
        None => Ok(None),
        Some(raw) => raw.parse::<usize>().map(Some).map_err(|err| {
            FormulaDslError::InvalidArgument {
                reason: format!(
                    "option `{key}={raw}` is not a non-negative integer; \
                     expected a whole number >= 0: {err}"
                ),
            }
            .into()
        }),
    }
}

/// Strict variant of `option_usize_any` that errors on the first present-but-
/// unparseable key rather than silently falling through.
pub fn option_usize_any_strict(
    map: &BTreeMap<String, String>,
    keys: &[&str],
) -> Result<Option<usize>, String> {
    for key in keys {
        if let Some(v) = option_usize_strict(map, key)? {
            return Ok(Some(v));
        }
    }
    Ok(None)
}

pub fn option_f64(map: &BTreeMap<String, String>, key: &str) -> Option<f64> {
    map.get(key).and_then(|v| v.parse::<f64>().ok())
}

/// Strict float option: `Ok(None)` if absent, `Ok(Some(n))` if parses as a
/// finite f64, `Err` if the user passed an unparseable value (rather than
/// silently dropping it like the lenient `option_f64`).
pub fn option_f64_strict(map: &BTreeMap<String, String>, key: &str) -> Result<Option<f64>, String> {
    match map.get(key) {
        None => Ok(None),
        Some(raw) => match raw.parse::<f64>() {
            Ok(v) if v.is_finite() => Ok(Some(v)),
            Ok(v) => Err(FormulaDslError::InvalidArgument {
                reason: format!("option `{key}={raw}` parses as {v} which is not a finite number"),
            }
            .into()),
            Err(err) => Err(FormulaDslError::InvalidArgument {
                reason: format!(
                    "option `{key}={raw}` is not a valid number; expected a finite decimal: {err}"
                ),
            }
            .into()),
        },
    }
}

pub fn option_bool(map: &BTreeMap<String, String>, key: &str) -> Option<bool> {
    map.get(key)
        .and_then(|v| match v.trim().to_ascii_lowercase().as_str() {
            "true" | "1" | "yes" | "y" => Some(true),
            "false" | "0" | "no" | "n" => Some(false),
            _ => None,
        })
}

/// Strict boolean option: `Ok(None)` if absent, `Ok(Some(b))` for a recognized
/// truthy/falsy token, and `Err(msg)` for a present-but-unparseable value. The
/// lenient `option_bool` maps an unrecognized value to `None`, which callers
/// then silently treat as "not specified" — masking user typos like
/// `double_penalty=ture`.
pub fn option_bool_strict(
    map: &BTreeMap<String, String>,
    key: &str,
) -> Result<Option<bool>, String> {
    match map.get(key) {
        None => Ok(None),
        Some(raw) => match raw.trim().to_ascii_lowercase().as_str() {
            "true" | "1" | "yes" | "y" => Ok(Some(true)),
            "false" | "0" | "no" | "n" => Ok(Some(false)),
            _ => Err(FormulaDslError::InvalidArgument {
                reason: format!(
                    "option `{key}={raw}` is not a boolean; \
                     expected one of true/false/yes/no/1/0"
                ),
            }
            .into()),
        },
    }
}

pub fn strip_quotes(v: &str) -> &str {
    let b = v.as_bytes();
    if b.len() >= 2
        && ((b[0] == b'\'' && b[b.len() - 1] == b'\'') || (b[0] == b'"' && b[b.len() - 1] == b'"'))
    {
        &v[1..v.len() - 1]
    } else {
        v
    }
}

// ---------------------------------------------------------------------------
// Sub-parsers for formula option blocks
// ---------------------------------------------------------------------------

fn parse_linear_constraint_bounds(
    options: &BTreeMap<String, String>,
    raw: &str,
) -> Result<(Option<f64>, Option<f64>), String> {
    let min = parse_optional_f64_option_alias(options, &["min", "lower"], raw, "linear")?;
    let max = parse_optional_f64_option_alias(options, &["max", "upper"], raw, "linear")?;
    if let (Some(min), Some(max)) = (min, max)
        && (!min.is_finite() || !max.is_finite() || min > max)
    {
        return Err(FormulaDslError::InvalidArgument {
            reason: format!(
                "linear coefficient constraints require finite min <= max, got min={min}, max={max}: {raw}"
            ),
        }
        .into());
    }
    Ok((min, max))
}

fn parse_required_f64_option(
    options: &BTreeMap<String, String>,
    key: &str,
    raw: &str,
) -> Result<f64, String> {
    let value = options
        .get(key)
        .ok_or_else(|| FormulaDslError::MalformedConfig {
            reason: format!("bounded() is missing required '{key}' argument: {raw}"),
        })?;
    value.parse::<f64>().map_err(|err| {
        FormulaDslError::InvalidArgument {
            reason: format!(
                "bounded() argument '{key}' must be a finite number, got '{}': {err}: {raw}",
                value
            ),
        }
        .into()
    })
}

fn parse_optional_f64_option(
    options: &BTreeMap<String, String>,
    key: &str,
    raw: &str,
) -> Result<Option<f64>, String> {
    match options.get(key) {
        Some(value) => value.parse::<f64>().map(Some).map_err(|err| {
            FormulaDslError::InvalidArgument {
                reason: format!(
                    "bounded() argument '{key}' must be a finite number, got '{}': {err}: {raw}",
                    value
                ),
            }
            .into()
        }),
        None => Ok(None),
    }
}

fn parse_optional_f64_option_alias(
    options: &BTreeMap<String, String>,
    keys: &[&str],
    raw: &str,
    fn_label: &str,
) -> Result<Option<f64>, String> {
    let mut found: Option<(&str, f64)> = None;
    for key in keys {
        if let Some(value) = options.get(*key) {
            let parsed = value
                .parse::<f64>()
                .map_err(|err| FormulaDslError::InvalidArgument {
                    reason: format!(
                        "{fn_label}() argument '{key}' must be a finite number, got '{}': {err}: {raw}",
                        value
                    ),
                })?;
            if found.is_some() {
                return Err(FormulaDslError::IncompatibleTerm {
                    reason: format!(
                        "{fn_label}() cannot specify both '{}' and '{}': {raw}",
                        found.expect("present").0,
                        key
                    ),
                }
                .into());
            }
            found = Some((key, parsed));
        }
    }
    Ok(found.map(|(_, v)| v))
}

fn parse_linkwiggle_penalty_orders(raw: Option<&str>) -> Result<Vec<usize>, String> {
    let Some(raw) = raw.map(str::trim) else {
        return Ok(WigglePenaltyConfig::cubic_triple_operator_default().penalty_orders);
    };
    if raw.is_empty() {
        return Ok(WigglePenaltyConfig::cubic_triple_operator_default().penalty_orders);
    }
    let mut out = Vec::<usize>::new();
    for token in raw.split(',') {
        let t = token.trim().to_ascii_lowercase();
        if t.is_empty() {
            continue;
        }
        match t.as_str() {
            "all" => {
                out.extend([1, 2, 3]);
            }
            "slope" | "1" => out.push(1),
            "curvature" | "2" => out.push(2),
            "curvature-change" | "curvature_change" | "3" => out.push(3),
            _ => {
                return Err(FormulaDslError::InvalidArgument {
                    reason: format!(
                        "invalid linkwiggle penalty_order '{t}'; use all|slope|curvature|curvature-change or 1/2/3"
                    ),
                }
                .into());
            }
        }
    }
    if out.is_empty() {
        out.extend(WigglePenaltyConfig::cubic_triple_operator_default().penalty_orders);
    }
    out.sort_unstable();
    out.dedup();
    Ok(out)
}

pub fn parse_linkwiggle_formulaspec(
    options: &BTreeMap<String, String>,
    raw: &str,
) -> Result<LinkWiggleFormulaSpec, String> {
    let allowed = [
        "degree",
        "internal_knots",
        "penalty_order",
        "double_penalty",
    ];
    let unknown = options
        .keys()
        .filter(|key| !allowed.contains(&key.as_str()))
        .cloned()
        .collect::<Vec<_>>();
    let term_name = raw.split('(').next().unwrap_or("linkwiggle");
    if !unknown.is_empty() {
        return Err(FormulaDslError::InvalidArgument {
            reason: format!(
                "{}() does not support option(s) {}: {raw}",
                term_name,
                unknown.join(", ")
            ),
        }
        .into());
    }
    let defaults = WigglePenaltyConfig::cubic_triple_operator_default();
    // Strict parsing: a present-but-unparseable value (`degree=abc`, `=-3`,
    // `=6.5`) must be rejected, not silently dropped and replaced by the
    // default as the lossy `option_usize`/`option_bool` readers would do.
    //
    // This parser is shared by *all* wiggle grammars: `linkwiggle` and
    // `timewiggle` (see `parse_formula`), the standard-model flexible-link
    // wiggle, and the marginal-slope score-warp / link-deviation routing.
    // The general monotone I-spline value basis (`monotone_wiggle_*` in
    // `families::gamlss`, used by `timewiggle` and the location-scale survival
    // path) honors arbitrary `degree >= 2`, while only the cubic-only
    // score-warp / link-deviation `DeviationRuntime` is restricted to 3.
    // A consumer-specific limit therefore must NOT be baked into this shared
    // parser — it is enforced at the routing layer that feeds the cubic-only
    // runtime (`deviation_block_config_from_formula_linkwiggle`). Here we only
    // enforce the universal lower bound that a polynomial degree is positive.
    let degree = option_usize_strict(options, "degree")?.unwrap_or(defaults.degree);
    if degree < 1 {
        return Err(FormulaDslError::InvalidArgument {
            reason: format!("{term_name}() requires degree >= 1: {raw}"),
        }
        .into());
    }
    let num_internal_knots =
        option_usize_strict(options, "internal_knots")?.unwrap_or(defaults.num_internal_knots);
    if num_internal_knots == 0 {
        return Err(FormulaDslError::InvalidArgument {
            reason: format!("{term_name}() requires internal_knots > 0: {raw}"),
        }
        .into());
    }
    let penalty_orders =
        parse_linkwiggle_penalty_orders(options.get("penalty_order").map(String::as_str))?;
    let double_penalty =
        option_bool_strict(options, "double_penalty")?.unwrap_or(defaults.double_penalty);
    Ok(LinkWiggleFormulaSpec {
        degree,
        num_internal_knots,
        penalty_orders,
        double_penalty,
    })
}

fn parse_link_formulaspec(
    options: &BTreeMap<String, String>,
    raw: &str,
) -> Result<LinkFormulaSpec, String> {
    let link = options
        .get("type")
        .map(|s| s.trim().to_string())
        .ok_or_else(|| FormulaDslError::MalformedConfig {
            reason: format!("link() requires type=<link-name>: {raw}"),
        })?;
    if link.is_empty() {
        return Err(FormulaDslError::MalformedConfig {
            reason: format!("link() requires a non-empty type: {raw}"),
        }
        .into());
    }
    let mixture_rho = options.get("rho").map(|s| s.trim().to_string());
    let sas_init = options.get("sas_init").map(|s| s.trim().to_string());
    let beta_logistic_init = options
        .get("beta_logistic_init")
        .map(|s| s.trim().to_string());
    Ok(LinkFormulaSpec {
        link,
        mixture_rho,
        sas_init,
        beta_logistic_init,
    })
}

fn parse_survival_formulaspec(
    options: &BTreeMap<String, String>,
    raw: &str,
) -> Result<SurvivalFormulaSpec, String> {
    if options.is_empty() {
        return Err(FormulaDslError::MalformedConfig {
            reason: format!(
                "survmodel() requires at least one named option (e.g., spec=..., distribution=...): {raw}"
            ),
        }
        .into());
    }
    Ok(SurvivalFormulaSpec {
        spec: options.get("spec").map(|s| s.trim().to_string()),
        survival_distribution: options.get("distribution").map(|s| s.trim().to_string()),
    })
}

fn parse_bounded_priorspec(
    options: &BTreeMap<String, String>,
    min: f64,
    max: f64,
    raw: &str,
) -> Result<BoundedCoefficientPriorSpec, String> {
    let prior_mode = options.get("prior").map(|s| s.to_ascii_lowercase());
    let pull = options.get("pull").map(|s| s.to_ascii_lowercase());
    let target = parse_optional_f64_option(options, "target", raw)?;
    let strength = parse_optional_f64_option(options, "strength", raw)?;

    let target_mode = target.is_some() || strength.is_some();
    if prior_mode.is_some() && pull.is_some() {
        return Err(FormulaDslError::IncompatibleTerm {
            reason: format!("bounded() cannot combine prior=... with pull=...: {raw}"),
        }
        .into());
    }
    if prior_mode.is_some() && target_mode {
        return Err(FormulaDslError::IncompatibleTerm {
            reason: format!("bounded() cannot combine prior=... with target/strength: {raw}"),
        }
        .into());
    }
    if pull.is_some() && target_mode {
        return Err(FormulaDslError::IncompatibleTerm {
            reason: format!("bounded() cannot combine pull=... with target/strength: {raw}"),
        }
        .into());
    }

    if let Some(priorname) = prior_mode {
        return match priorname.as_str() {
            "none" => Ok(BoundedCoefficientPriorSpec::None),
            "uniform" | "log-jacobian" | "log_jacobian" | "jacobian" => {
                Ok(BoundedCoefficientPriorSpec::Uniform)
            }
            "center" => Ok(BoundedCoefficientPriorSpec::Beta { a: 2.0, b: 2.0 }),
            _ => Err(FormulaDslError::InvalidArgument {
                reason: format!(
                    "bounded() prior must currently be one of none|uniform|log-jacobian|center, got '{}': {raw}",
                    priorname
                ),
            }
            .into()),
        };
    }

    if let Some(pull_mode) = pull {
        return match pull_mode.as_str() {
            "uniform" | "log-jacobian" | "log_jacobian" | "jacobian" => {
                Ok(BoundedCoefficientPriorSpec::Uniform)
            }
            "center" => Ok(BoundedCoefficientPriorSpec::Beta { a: 2.0, b: 2.0 }),
            _ => Err(FormulaDslError::InvalidArgument {
                reason: format!(
                    "bounded() pull must currently be 'uniform'/'log-jacobian' or 'center', got '{}': {raw}",
                    pull_mode
                ),
            }
            .into()),
        };
    }

    if target_mode {
        let targetvalue = target.ok_or_else(|| FormulaDslError::MalformedConfig {
            reason: format!("bounded() target is required with strength: {raw}"),
        })?;
        let strengthvalue = strength.ok_or_else(|| FormulaDslError::MalformedConfig {
            reason: format!("bounded() strength is required with target: {raw}"),
        })?;
        if !(min < targetvalue && targetvalue < max) {
            return Err(FormulaDslError::InvalidArgument {
                reason: format!("bounded() target must lie strictly inside ({min}, {max}): {raw}"),
            }
            .into());
        }
        if !strengthvalue.is_finite() || strengthvalue <= 0.0 {
            return Err(FormulaDslError::InvalidArgument {
                reason: format!("bounded() strength must be finite and > 0: {raw}"),
            }
            .into());
        }
        let z = (targetvalue - min) / (max - min);
        let a = 1.0 + strengthvalue * z;
        let b = 1.0 + strengthvalue * (1.0 - z);
        return Ok(BoundedCoefficientPriorSpec::Beta { a, b });
    }

    Ok(BoundedCoefficientPriorSpec::None)
}

// ---------------------------------------------------------------------------
// Top-level formula and term parsers
// ---------------------------------------------------------------------------

pub fn formula_rhs_text(formula: &str) -> Result<String, String> {
    let parsed = parse_formula_dsl(formula)?;
    if parsed.rhs_terms.is_empty() {
        return Err(FormulaDslError::ParseError {
            reason: "formula right-hand side cannot be empty".to_string(),
        }
        .into());
    }
    Ok(parsed.rhs_terms.join(" + "))
}

/// Parsed Surv(...) response specification.
///
/// `entry` is `None` for the 2-arg right-censored shorthand
/// `Surv(time, event)`, which matches the R survival/mgcv default: every
/// subject has entry time zero. Callers materialize a zero entry column
/// when this is `None`.
pub fn parse_surv_response(
    lhs: &str,
) -> Result<Option<(Option<String>, String, String)>, FormulaDslError> {
    let trimmed = lhs.trim();
    let call = match parse_function_call(trimmed) {
        Ok(call) => call,
        Err(_) => return Ok(None),
    };
    if !call.name.eq_ignore_ascii_case("surv") {
        return Ok(None);
    }
    let vars = call
        .args
        .iter()
        .filter_map(|arg| match arg {
            CallArgSpec::Positional(v) => Some(v.trim().to_string()),
            CallArgSpec::Named { .. } => None,
        })
        .filter(|s| !s.is_empty())
        .collect::<Vec<_>>();
    match vars.len() {
        // Right-censored shorthand: Surv(time, event) ≡ Surv(0, time, event)
        // with a synthetic zero entry column. This matches R's
        // `survival::Surv(time, event)` default for left-truncation-free data.
        2 => Ok(Some((None, vars[0].clone(), vars[1].clone()))),
        3 => Ok(Some((
            Some(vars[0].clone()),
            vars[1].clone(),
            vars[2].clone(),
        ))),
        n => Err(FormulaDslError::InvalidArgument {
            reason: format!(
                "Surv(...) expects either Surv(time, event) (right-censored) or \
                 Surv(entry, exit, event) (left-truncated); got {n} columns"
            ),
        }),
    }
}

/// Parsed `SurvInterval(L, R, event)` interval-censored response.
///
/// Returns `Some((left_col, right_col, event_col))` when the left-hand side is a
/// `SurvInterval(...)` call, `None` otherwise (so a plain `Surv(...)` or a bare
/// column response falls through to the other response parsers).
///
/// Interval censoring observes only a bracket `T ∈ (L, R]` — the exact event
/// time is never seen — and its row contribution is the survival-mass difference
/// `log[S(L) − S(R)]`, distinct from both the exact-event point density and the
/// single-sided right-censored survival. A *dedicated call name* (rather than
/// overloading the 3-argument `Surv(entry, exit, event)` delayed-entry form,
/// which is also 3-argument and semantically incompatible) is the unambiguous
/// DSL spelling: it mirrors flexsurv's `Surv(L, R, type="interval2")` intent
/// without colliding with the existing left-truncation grammar.
pub fn parse_surv_interval_response(
    lhs: &str,
) -> Result<Option<(String, String, String)>, FormulaDslError> {
    let trimmed = lhs.trim();
    let call = match parse_function_call(trimmed) {
        Ok(call) => call,
        Err(_) => return Ok(None),
    };
    if !call.name.eq_ignore_ascii_case("survinterval") {
        return Ok(None);
    }
    let vars = call
        .args
        .iter()
        .filter_map(|arg| match arg {
            CallArgSpec::Positional(v) => Some(v.trim().to_string()),
            CallArgSpec::Named { .. } => None,
        })
        .filter(|s| !s.is_empty())
        .collect::<Vec<_>>();
    match vars.len() {
        3 => Ok(Some((vars[0].clone(), vars[1].clone(), vars[2].clone()))),
        n => Err(FormulaDslError::InvalidArgument {
            reason: format!(
                "SurvInterval(...) expects SurvInterval(L, R, event) (interval-censored, the \
                 observed bracket T ∈ (L, R]); got {n} columns"
            ),
        }),
    }
}

fn top_level_formula_separator(input: &str) -> Result<Option<usize>, String> {
    let mut depth = 0_i32;
    let mut in_single = false;
    let mut in_double = false;

    for (idx, ch) in input.char_indices() {
        match ch {
            '\'' if !in_double => in_single = !in_single,
            '"' if !in_single => in_double = !in_double,
            '(' | '[' | '{' if !in_single && !in_double => depth += 1,
            ')' | ']' | '}' if !in_single && !in_double && depth > 0 => depth -= 1,
            '~' if !in_single && !in_double && depth == 0 => return Ok(Some(idx)),
            _ => {}
        }
    }

    if in_single || in_double || depth != 0 {
        return Err(FormulaDslError::ParseError {
            reason: "invalid auxiliary formula syntax: unbalanced parentheses or quotes"
                .to_string(),
        }
        .into());
    }
    Ok(None)
}

pub fn parse_matching_auxiliary_formula(
    formula: &str,
    response: &str,
    flag_name: &str,
) -> Result<(String, ParsedFormula), FormulaDslError> {
    let rhs = formula.trim();
    if top_level_formula_separator(rhs)?.is_some() {
        return Err(FormulaDslError::InvalidArgument {
            reason: format!(
                "{flag_name} expects only the terms after '~', not a full 'response ~ terms' formula; use {flag_name} 's(x)' instead of {flag_name} 'y ~ s(x)' (or pass '1' for an intercept-only noise model)"
            ),
        });
    }
    let parsed_formula = parse_formula(&format!("{response} ~ {rhs}"))?;
    Ok((rhs.to_string(), parsed_formula))
}

pub fn validate_auxiliary_formula_controls(
    parsed_formula: &ParsedFormula,
    flag_name: &str,
) -> Result<(), String> {
    if parsed_formula.linkwiggle.is_some() {
        return Err(FormulaDslError::IncompatibleTerm {
            reason: format!(
                "linkwiggle(...) is only supported in the main formula, not {flag_name}"
            ),
        }
        .into());
    }
    if parsed_formula.timewiggle.is_some() {
        return Err(FormulaDslError::IncompatibleTerm {
            reason: format!(
                "timewiggle(...) is only supported in the main survival formula, not {flag_name}"
            ),
        }
        .into());
    }
    if parsed_formula.linkspec.is_some() {
        return Err(FormulaDslError::IncompatibleTerm {
            reason: format!("link(...) is only supported in the main formula, not {flag_name}"),
        }
        .into());
    }
    if parsed_formula.survivalspec.is_some() {
        return Err(FormulaDslError::IncompatibleTerm {
            reason: format!(
                "survmodel(...) is only supported in the main survival formula, not {flag_name}"
            ),
        }
        .into());
    }
    if !parsed_formula.logslope_surfaces.is_empty() && flag_name != "--logslope-formula" {
        return Err(FormulaDslError::IncompatibleTerm {
            reason: format!(
                "logslope(...) is only supported in --logslope-formula, not {flag_name}"
            ),
        }
        .into());
    }
    Ok::<(), _>(())
}

pub fn parse_formula(formula: &str) -> Result<ParsedFormula, FormulaDslError> {
    let parsed_dsl =
        parse_formula_dsl(formula).map_err(|reason| FormulaDslError::ParseError { reason })?;
    let lhs = parsed_dsl.response_expr.trim();
    if lhs.is_empty() {
        return Err(FormulaDslError::ParseError {
            reason: "formula response (left-hand side) cannot be empty".to_string(),
        });
    }
    let mut terms = Vec::<ParsedTerm>::new();
    let mut linkwiggle: Option<LinkWiggleFormulaSpec> = None;
    let mut timewiggle: Option<LinkWiggleFormulaSpec> = None;
    let mut linkspec: Option<LinkFormulaSpec> = None;
    let mut survivalspec: Option<SurvivalFormulaSpec> = None;
    let mut logslope_surfaces = Vec::<LogSlopeSurfaceSpec>::new();
    // Track seen-term-keys so we can reject exact duplicates like
    // `y ~ smooth(x) + smooth(x)` upfront — without this the duplicate
    // produces a rank-deficient design and the user has no idea why their
    // fit is over-parameterized.
    let mut seen_term_keys: std::collections::BTreeSet<String> = std::collections::BTreeSet::new();
    let mut expanded_terms = Vec::<String>::new();
    for raw in parsed_dsl.rhs_terms {
        let trimmed = raw.trim();
        if trimmed.is_empty() {
            expanded_terms.push(String::new());
            continue;
        }
        // Single function-call terms (smooths, group(), etc.) are opaque to
        // WR expansion; pass them through verbatim so parse_term sees the
        // exact source string. Bare identifiers and any term mentioning
        // `:`, `*`, `/`, `^` are routed through the AST-driven WR expander.
        let is_call = parse_function_call(trimmed).is_ok();
        let needs_expansion = !is_call
            && trimmed
                .chars()
                .scan(0i32, |depth, ch| {
                    let d_before = *depth;
                    match ch {
                        '(' | '[' | '{' => *depth += 1,
                        ')' | ']' | '}' if *depth > 0 => *depth -= 1,
                        _ => {}
                    }
                    Some((d_before, ch))
                })
                .any(|(d, ch)| d == 0 && matches!(ch, ':' | '*' | '/' | '^'));
        if needs_expansion {
            for atoms in
                expand_wr_term(trimmed).map_err(|reason| FormulaDslError::ParseError { reason })?
            {
                if atoms.is_empty() {
                    continue;
                }
                expanded_terms.push(atoms.join(":"));
            }
        } else {
            expanded_terms.push(trimmed.to_string());
        }
    }

    for raw in expanded_terms {
        let t = raw.trim();
        if t.is_empty() || t == "1" {
            continue;
        }
        if t == "0" || t == "-1" {
            return Err(FormulaDslError::IncompatibleTerm {
                reason: "formula terms '0'/'-1' (intercept removal) are not supported yet"
                    .to_string(),
            });
        }
        // Normalize whitespace so `smooth(x)` and `smooth( x )` match,
        // but preserve whitespace inside string literals so that
        // `bs="a b"` and `bs="ab"` do not collide.
        let key: String = {
            let mut acc = String::with_capacity(t.len());
            let mut in_single = false;
            let mut in_double = false;
            for ch in t.chars() {
                match ch {
                    '\'' if !in_double => {
                        in_single = !in_single;
                        acc.push(ch);
                    }
                    '"' if !in_single => {
                        in_double = !in_double;
                        acc.push(ch);
                    }
                    c if c.is_whitespace() && !in_single && !in_double => {}
                    _ => acc.push(ch),
                }
            }
            acc
        };
        if !seen_term_keys.insert(key.clone()) {
            return Err(FormulaDslError::IncompatibleTerm {
                reason: format!(
                    "formula `{formula}` lists term `{t}` more than once. \
                     Duplicate terms produce a rank-deficient design; \
                     keep one copy or differentiate them (e.g. distinct k=, bs= options)."
                ),
            });
        }
        match parse_term(t)? {
            ParsedTerm::LinkWiggle { options } => {
                if linkwiggle.is_some() {
                    return Err(FormulaDslError::IncompatibleTerm {
                        reason: "formula can include at most one linkwiggle(...) term".to_string(),
                    });
                }
                linkwiggle = Some(parse_linkwiggle_formulaspec(&options, t)?);
            }
            ParsedTerm::TimeWiggle { options } => {
                if timewiggle.is_some() {
                    return Err(FormulaDslError::IncompatibleTerm {
                        reason: "formula can include at most one timewiggle(...) term".to_string(),
                    });
                }
                timewiggle = Some(parse_linkwiggle_formulaspec(&options, t)?);
            }
            ParsedTerm::LinkConfig { options } => {
                if linkspec.is_some() {
                    return Err(FormulaDslError::IncompatibleTerm {
                        reason: "formula can include at most one link(...) term".to_string(),
                    });
                }
                linkspec = Some(parse_link_formulaspec(&options, t)?);
            }
            ParsedTerm::SurvivalConfig { options } => {
                if survivalspec.is_some() {
                    return Err(FormulaDslError::IncompatibleTerm {
                        reason: "formula can include at most one survmodel(...) term".to_string(),
                    });
                }
                survivalspec = Some(parse_survival_formulaspec(&options, t)?);
            }
            ParsedTerm::LogSlopeSurface { z_column, terms } => {
                logslope_surfaces.push(LogSlopeSurfaceSpec { z_column, terms });
            }
            other => terms.push(other),
        }
    }
    // Reject self-referential formulas like `y ~ smooth(y)` or `y ~ y`: the
    // response is its own predictor, which is a trivial identity fit and
    // almost certainly a user mistake. Only flag the simple-identifier case
    // (so Surv(entry, exit, event) ~ smooth(time) is left alone — the
    // response is the Surv triple, not the bare "time" column).
    if lhs.chars().all(|c| c.is_alphanumeric() || c == '_')
        && !lhs.is_empty()
        && parsed_terms_reference_column(&terms, lhs)
    {
        return Err(FormulaDslError::IncompatibleTerm {
            reason: format!(
                "formula `{formula}` uses response column `{lhs}` as its own predictor. \
                 This fits y as a function of itself and is almost certainly a typo. \
                 Drop the term that mentions `{lhs}` from the right-hand side."
            ),
        });
    }
    Ok(ParsedFormula {
        response: lhs.to_string(),
        terms,
        logslope_surfaces,
        linkwiggle,
        timewiggle,
        linkspec,
        survivalspec,
    })
}

pub fn parse_term(raw: &str) -> Result<ParsedTerm, String> {
    fn split_call_args(call: &FunctionCallSpec) -> (Vec<String>, BTreeMap<String, String>) {
        let mut vars = Vec::<String>::new();
        let mut options = BTreeMap::<String, String>::new();
        for arg in &call.args {
            match arg {
                CallArgSpec::Positional(v) => vars.push(v.trim().to_string()),
                CallArgSpec::Named { key, value } => {
                    options.insert(key.to_ascii_lowercase(), strip_quotes(value).to_string());
                }
            }
        }
        (vars, options)
    }

    // Wilkinson-Rogers `:` interaction term. The expander in `parse_formula`
    // produces `a:b[:c...]` for these; parse_term is also reached directly
    // from tests, so handle the syntax here as well.
    if raw.contains(':')
        && !raw.contains('(')
        && raw.split(':').all(|piece| is_exact_ident(piece.trim()))
    {
        let vars: Vec<String> = raw
            .split(':')
            .map(|piece| piece.trim().to_string())
            .collect();
        if vars.len() >= 2 {
            let mut sorted = vars.clone();
            sorted.sort();
            sorted.dedup();
            if sorted.len() != vars.len() {
                return Err(FormulaDslError::IncompatibleTerm {
                    reason: format!(
                        "interaction term `{raw}` references the same variable more than once"
                    ),
                }
                .into());
            }
            return Ok(ParsedTerm::Interaction {
                vars: sorted,
                double_penalty: true,
            });
        }
    }

    let call = parse_function_call(raw).ok();
    if let Some(call) = call {
        let name = call.name.to_ascii_lowercase();
        let (vars, mut options) = split_call_args(&call);
        match name.as_str() {
            "constrain" | "constraint" | "box" => {
                if vars.len() != 1 {
                    return Err(FormulaDslError::InvalidArgument {
                        reason: format!(
                            "constrain()/constraint()/box() expects exactly one variable: {raw}"
                        ),
                    }
                    .into());
                }
                validate_known_term_options(
                    "constrain",
                    &options,
                    &["min", "lower", "max", "upper", "double_penalty"],
                    raw,
                )?;
                let (coefficient_min, coefficient_max) =
                    parse_linear_constraint_bounds(&options, raw)?;
                if coefficient_min.is_none() && coefficient_max.is_none() {
                    return Err(FormulaDslError::MalformedConfig {
                        reason: format!(
                            "constrain()/constraint()/box() requires at least one of min/lower/max/upper: {raw}"
                        ),
                    }
                    .into());
                }
                return Ok(ParsedTerm::Linear {
                    name: vars[0].clone(),
                    explicit: true,
                    double_penalty: option_bool_strict(&options, "double_penalty")?.unwrap_or(true),
                    coefficient_min,
                    coefficient_max,
                });
            }
            "nonnegative" | "nonnegative_coef" => {
                if vars.len() != 1 {
                    return Err(FormulaDslError::InvalidArgument {
                        reason: format!("nonnegative() expects exactly one variable: {raw}"),
                    }
                    .into());
                }
                validate_known_term_options("nonnegative", &options, &["double_penalty"], raw)?;
                return Ok(ParsedTerm::Linear {
                    name: vars[0].clone(),
                    explicit: true,
                    double_penalty: option_bool_strict(&options, "double_penalty")?.unwrap_or(true),
                    coefficient_min: Some(0.0),
                    coefficient_max: None,
                });
            }
            "nonpositive" | "nonpositive_coef" => {
                if vars.len() != 1 {
                    return Err(FormulaDslError::InvalidArgument {
                        reason: format!("nonpositive() expects exactly one variable: {raw}"),
                    }
                    .into());
                }
                validate_known_term_options("nonpositive", &options, &["double_penalty"], raw)?;
                return Ok(ParsedTerm::Linear {
                    name: vars[0].clone(),
                    explicit: true,
                    double_penalty: option_bool_strict(&options, "double_penalty")?.unwrap_or(true),
                    coefficient_min: None,
                    coefficient_max: Some(0.0),
                });
            }
            "bounded" => {
                if vars.len() != 1 {
                    return Err(FormulaDslError::InvalidArgument {
                        reason: format!("bounded() expects exactly one variable: {raw}"),
                    }
                    .into());
                }
                validate_known_term_options(
                    "bounded",
                    &options,
                    &[
                        "min",
                        "max",
                        "prior",
                        "pull",
                        "target",
                        "strength",
                        "double_penalty",
                    ],
                    raw,
                )?;
                let min = parse_required_f64_option(&options, "min", raw)?;
                let max = parse_required_f64_option(&options, "max", raw)?;
                if !min.is_finite() || !max.is_finite() || min >= max {
                    return Err(FormulaDslError::InvalidArgument {
                        reason: format!(
                            "bounded() requires finite min < max, got min={min}, max={max}: {raw}"
                        ),
                    }
                    .into());
                }
                let prior = parse_bounded_priorspec(&options, min, max, raw)?;
                return Ok(ParsedTerm::BoundedLinear {
                    name: vars[0].clone(),
                    min,
                    max,
                    prior,
                    // Unlike a plain `linear()` term, `bounded()` already commits
                    // the coefficient to an exact interval transform (plus an
                    // optional prior); layering the null-space ridge on top is
                    // structurally rejected downstream (`design_construction.rs`:
                    // "bounded linear term ... cannot also use double_penalty"),
                    // so the default must be `false`, not the `linear()`/`s()`
                    // convention of `true`.
                    double_penalty: option_bool_strict(&options, "double_penalty")?
                        .unwrap_or(false),
                });
            }
            "group" | "re" | "factor" => {
                if vars.len() != 1 {
                    return Err(FormulaDslError::InvalidArgument {
                        reason: format!(
                            "{name}() expects exactly one variable, got '{}': {raw}",
                            vars.join(",")
                        ),
                    }
                    .into());
                }
                // `factor(g)` is a FIXED categorical factor (R `factor()` /
                // patsy `C()`): it forces categorical encoding of the column
                // but, like a bare `+ g` main effect, is strict on unseen
                // levels. `group(g)`/`re(g)` are genuine random effects that
                // shrink a held-out group to the population mean, so they
                // tolerate unseen levels. Both share the penalized-categorical
                // block; only the unseen policy differs (#2137/#2102).
                let lenient_unseen = name != "factor";
                return Ok(ParsedTerm::RandomEffect {
                    name: vars[0].clone(),
                    lenient_unseen,
                });
            }
            "tensor" | "interaction" | "te" => {
                if vars.len() < 2 {
                    return Err(FormulaDslError::InvalidArgument {
                        reason: format!(
                            "tensor()/interaction()/te() requires at least two variables: {raw}"
                        ),
                    }
                    .into());
                }
                return Ok(ParsedTerm::Smooth {
                    label: raw.to_string(),
                    vars,
                    kind: SmoothKind::Te,
                    options,
                });
            }
            "t2" => {
                if vars.len() < 2 {
                    return Err(FormulaDslError::InvalidArgument {
                        reason: format!("t2() requires at least two variables: {raw}"),
                    }
                    .into());
                }
                return Ok(ParsedTerm::Smooth {
                    label: raw.to_string(),
                    vars,
                    kind: SmoothKind::T2,
                    options,
                });
            }
            "ti" => {
                // Tensor interaction smooth (mgcv `ti`): structurally a
                // tensor-product smooth, but the marginal main effects are
                // excluded so only the pure interaction is modeled. Shares the
                // tensor materialization path with `te`; the distinct
                // `SmoothKind::Ti` drives per-margin sum-to-zero
                // identifiability in the term builder.
                if vars.len() < 2 {
                    return Err(FormulaDslError::InvalidArgument {
                        reason: format!("ti() requires at least two variables: {raw}"),
                    }
                    .into());
                }
                return Ok(ParsedTerm::Smooth {
                    label: raw.to_string(),
                    vars,
                    kind: SmoothKind::Ti,
                    options,
                });
            }
            "fs" | "sz" => {
                if vars.len() != 2 {
                    return Err(format!("{}() expects exactly two variables: {raw}", name));
                }
                options.insert("bs".to_string(), name.clone());
                return Ok(ParsedTerm::Smooth {
                    label: raw.to_string(),
                    vars,
                    kind: SmoothKind::S,
                    options,
                });
            }
            "thinplate" | "thin_plate" | "tps" => {
                if vars.len() < 2 {
                    return Err(FormulaDslError::InvalidArgument {
                        reason: format!(
                            "thinplate()/thin_plate()/tps() requires at least two variables: {raw}"
                        ),
                    }
                    .into());
                }
                options.insert("type".to_string(), "tps".to_string());
                return Ok(ParsedTerm::Smooth {
                    label: raw.to_string(),
                    vars,
                    kind: SmoothKind::S,
                    options,
                });
            }
            "smooth" | "s" | "cyclic" | "periodic" | "cc" | "cp" => {
                if vars.is_empty() {
                    return Err(FormulaDslError::InvalidArgument {
                        reason: format!("smooth()/s() requires at least one variable: {raw}"),
                    }
                    .into());
                }
                // mgcv idiom: `s(g, bs='re')` with a single variable is a
                // random intercept on the factor `g`. Route it to the
                // dedicated random-effect machinery (which expects a single
                // categorical column) rather than to the factor-smooth path
                // (which requires a numeric companion).
                let bs_is_re = options
                    .get("bs")
                    .or_else(|| options.get("type"))
                    .map(|v| {
                        v.trim()
                            .trim_matches(|c| c == '\'' || c == '"')
                            .to_ascii_lowercase()
                    })
                    .as_deref()
                    == Some("re");
                if bs_is_re && vars.len() == 1 {
                    // `s(g, bs="re")` is a genuine random effect: lenient on
                    // unseen levels (held-out group → population mean).
                    return Ok(ParsedTerm::RandomEffect {
                        name: vars[0].clone(),
                        lenient_unseen: true,
                    });
                }
                if matches!(name.as_str(), "cyclic" | "periodic" | "cc" | "cp") {
                    options.insert("type".to_string(), "cyclic".to_string());
                }
                if matches!(name.as_str(), "fs" | "sz") {
                    options.insert("bs".to_string(), name.clone());
                }
                return Ok(ParsedTerm::Smooth {
                    label: raw.to_string(),
                    vars,
                    kind: SmoothKind::S,
                    options,
                });
            }
            "sphere" | "sos" | "spherical" | "s2" => {
                // `s2()` is an alias for the intrinsic S² (sphere) smooth, just
                // like `sphere()`/`sos()`/`spherical()`. All four share the
                // Wahba/harmonic sphere basis, so they must dispatch through
                // the identical `type=sphere` route; otherwise `s2()` would
                // silently fall back to a generic Euclidean 2-D smooth over
                // (lat, lon) and diverge in the spatial-kappa optimizer.
                if vars.len() != 2 {
                    return Err(FormulaDslError::InvalidArgument {
                        reason: format!(
                            "{name}() expects exactly two variables: latitude and longitude; got {} in {raw}",
                            vars.len()
                        ),
                    }
                    .into());
                }
                options.insert("type".to_string(), "sphere".to_string());
                return Ok(ParsedTerm::Smooth {
                    label: raw.to_string(),
                    vars,
                    kind: SmoothKind::S,
                    options,
                });
            }
            "mjs" | "measurejet" | "measure_jet" | "web" => {
                // Measure-jet spline smooth (`basis::measure_jet_smooth` docs)
                // for responses varying along an unknown low-dimensional set
                // inside a higher-dimensional ambient space. All aliases
                // dispatch through the identical `type=measurejet` route,
                // mirroring the sphere/curvature alias rule.
                if vars.is_empty() {
                    return Err(FormulaDslError::InvalidArgument {
                        reason: format!("{name}() requires at least one variable: {raw}"),
                    }
                    .into());
                }
                options.insert("type".to_string(), "measurejet".to_string());
                return Ok(ParsedTerm::Smooth {
                    label: raw.to_string(),
                    vars,
                    kind: SmoothKind::S,
                    options,
                });
            }
            "curv" | "curvature" | "constant_curvature" | "mkappa" => {
                // Constant-curvature (M_κ) geodesic-kernel smooth (#944): the
                // κ-generic sibling of sphere()/s2(), interpolating
                // S^d → ℝ^d → H^d through `kappa=` (default 0 = flat). All
                // four aliases must dispatch through the identical
                // `type=curvature` route, mirroring the sphere alias rule.
                if vars.is_empty() {
                    return Err(FormulaDslError::InvalidArgument {
                        reason: format!("{name}() requires at least one variable: {raw}"),
                    }
                    .into());
                }
                options.insert("type".to_string(), "curvature".to_string());
                return Ok(ParsedTerm::Smooth {
                    label: raw.to_string(),
                    vars,
                    kind: SmoothKind::S,
                    options,
                });
            }
            "matern" => {
                if vars.is_empty() {
                    return Err(FormulaDslError::InvalidArgument {
                        reason: format!("matern() requires at least one variable: {raw}"),
                    }
                    .into());
                }
                options.insert("type".to_string(), "matern".to_string());
                return Ok(ParsedTerm::Smooth {
                    label: raw.to_string(),
                    vars,
                    kind: SmoothKind::S,
                    options,
                });
            }
            "duchon" => {
                if vars.is_empty() {
                    return Err(FormulaDslError::InvalidArgument {
                        reason: format!("duchon() requires at least one variable: {raw}"),
                    }
                    .into());
                }
                if option_bool(&options, "cyclic").unwrap_or(false)
                    || option_bool(&options, "periodic").unwrap_or(false)
                {
                    options.insert("cyclic".to_string(), "true".to_string());
                }
                options.insert("type".to_string(), "duchon".to_string());
                return Ok(ParsedTerm::Smooth {
                    label: raw.to_string(),
                    vars,
                    kind: SmoothKind::S,
                    options,
                });
            }
            "pca" => {
                if vars.is_empty() {
                    return Err(FormulaDslError::InvalidArgument {
                        reason: format!("pca() requires at least one variable: {raw}"),
                    }
                    .into());
                }
                options.insert("type".to_string(), "pca".to_string());
                return Ok(ParsedTerm::Smooth {
                    label: raw.to_string(),
                    vars,
                    kind: SmoothKind::S,
                    options,
                });
            }
            "linkwiggle" => {
                if !vars.is_empty() {
                    return Err(FormulaDslError::InvalidArgument {
                        reason: format!(
                            "linkwiggle() takes named options only; positional args are not supported: {raw}"
                        ),
                    }
                    .into());
                }
                return Ok(ParsedTerm::LinkWiggle { options });
            }
            "timewiggle" => {
                if !vars.is_empty() {
                    return Err(FormulaDslError::InvalidArgument {
                        reason: format!(
                            "timewiggle() takes named options only; positional args are not supported: {raw}"
                        ),
                    }
                    .into());
                }
                return Ok(ParsedTerm::TimeWiggle { options });
            }
            "link" => {
                if !vars.is_empty() {
                    return Err(FormulaDslError::InvalidArgument {
                        reason: format!(
                            "link() takes named options only; positional args are not supported: {raw}"
                        ),
                    }
                    .into());
                }
                return Ok(ParsedTerm::LinkConfig { options });
            }
            "survmodel" => {
                if !vars.is_empty() {
                    return Err(FormulaDslError::InvalidArgument {
                        reason: format!(
                            "survmodel() takes named options only; positional args are not supported: {raw}"
                        ),
                    }
                    .into());
                }
                return Ok(ParsedTerm::SurvivalConfig { options });
            }
            "logslope" | "log_slope" | "log_slope_surface" => {
                validate_known_term_options("logslope", &options, &[], raw)?;
                if vars.len() < 2 {
                    return Err(FormulaDslError::InvalidArgument {
                        reason: format!(
                            "logslope() expects a z column followed by one or more RHS terms; add one logslope(z, ...) declaration per vector-z coordinate: {raw}"
                        ),
                    }
                    .into());
                }
                let z_column = vars[0].trim();
                if !is_exact_ident(z_column) {
                    return Err(FormulaDslError::InvalidArgument {
                        reason: format!(
                            "logslope() z column must be a bare column name, got `{z_column}` in {raw}"
                        ),
                    }
                    .into());
                }
                let rhs = vars[1..].join(" + ");
                let parsed = parse_formula(&format!("__logslope__ ~ {rhs}"))?;
                if !parsed.logslope_surfaces.is_empty() {
                    return Err(FormulaDslError::IncompatibleTerm {
                        reason: format!(
                            "logslope() declarations cannot be nested inside another logslope(): {raw}"
                        ),
                    }
                    .into());
                }
                validate_auxiliary_formula_controls(&parsed, "logslope()")?;
                return Ok(ParsedTerm::LogSlopeSurface {
                    z_column: z_column.to_string(),
                    terms: parsed.terms,
                });
            }
            "linear" => {
                if vars.len() != 1 {
                    return Err(FormulaDslError::InvalidArgument {
                        reason: format!("linear() expects exactly one variable: {raw}"),
                    }
                    .into());
                }
                validate_known_term_options(
                    "linear",
                    &options,
                    &["min", "lower", "max", "upper", "double_penalty"],
                    raw,
                )?;
                let (coefficient_min, coefficient_max) =
                    parse_linear_constraint_bounds(&options, raw)?;
                let double_penalty =
                    option_bool_strict(&options, "double_penalty")?.unwrap_or(true);
                if vars[0].contains(':') {
                    if coefficient_min.is_some() || coefficient_max.is_some() {
                        return Err(FormulaDslError::IncompatibleTerm {
                            reason: format!(
                                "linear() coefficient bounds are not supported on an interaction: {raw}"
                            ),
                        }
                        .into());
                    }
                    let mut interaction_vars = vars[0]
                        .split(':')
                        .map(str::trim)
                        .map(str::to_string)
                        .collect::<Vec<_>>();
                    if interaction_vars.len() < 2
                        || interaction_vars.iter().any(|var| !is_exact_ident(var))
                    {
                        return Err(FormulaDslError::InvalidArgument {
                            reason: format!(
                                "linear() interaction must contain at least two bare column names: {raw}"
                            ),
                        }
                        .into());
                    }
                    interaction_vars.sort();
                    let original_len = interaction_vars.len();
                    interaction_vars.dedup();
                    if interaction_vars.len() != original_len {
                        return Err(FormulaDslError::IncompatibleTerm {
                            reason: format!(
                                "linear() interaction references the same variable more than once: {raw}"
                            ),
                        }
                        .into());
                    }
                    return Ok(ParsedTerm::Interaction {
                        vars: interaction_vars,
                        double_penalty,
                    });
                }
                return Ok(ParsedTerm::Linear {
                    name: vars[0].clone(),
                    explicit: true,
                    double_penalty,
                    coefficient_min,
                    coefficient_max,
                });
            }
            _ => {
                return Err(format!(
                    "unknown term function `{name}` in '{raw}'. Supported: bounded(), linear(), constrain()/constraint()/box(), nonnegative(), nonpositive(), smooth()/s(), cyclic()/periodic()/cc()/cp(), thinplate()/thin_plate()/tps(), tensor()/interaction()/te(), t2(), ti(), fs(), sz(), group()/re()/factor(), sphere()/sos()/spherical(), s2(), matern(), duchon(), pca(), logslope()/log_slope(), linkwiggle(), timewiggle(), link(), survmodel()"
                ));
            }
        }
    }

    let ident = raw.trim();
    if !is_exact_ident(ident) {
        return Err(FormulaDslError::UnknownIdentifier {
            reason: format!("unsupported top-level RHS term: {raw}"),
        }
        .into());
    }

    Ok(ParsedTerm::Linear {
        name: ident.to_string(),
        explicit: false,
        double_penalty: true,
        coefficient_min: None,
        coefficient_max: None,
    })
}

// ---------------------------------------------------------------------------
// Link choice parsing
// ---------------------------------------------------------------------------

pub fn parse_link_choice(
    raw: Option<&str>,
    flexible_flag: bool,
) -> Result<Option<LinkChoice>, FormulaDslError> {
    if raw.is_none() && !flexible_flag {
        return Ok(None);
    }
    let Some(v) = raw else {
        return Ok(Some(LinkChoice {
            mode: LinkMode::Flexible,
            link: LinkFunction::Probit,
            mixture_components: None,
        }));
    };
    let t = v.trim().to_ascii_lowercase();
    if let Some(inner) = t
        .strip_prefix("flexible(")
        .and_then(|s| s.strip_suffix(')'))
    {
        if let Some(components_inner) = inner
            .strip_prefix("blended(")
            .and_then(|s| s.strip_suffix(')'))
            .or_else(|| {
                inner
                    .strip_prefix("mixture(")
                    .and_then(|s| s.strip_suffix(')'))
            })
        {
            parse_link_component_list(components_inner)?;
            return Err(FormulaDslError::IncompatibleTerm {
                reason:
                    "flexible(...) does not support blended(...)/mixture(...) links; wiggle is only supported for jointly fit standard links"
                        .to_string(),
            });
        }
        let link = parse_linkname(inner)?;
        if !linkname_supports_joint_wiggle(link) {
            return Err(FormulaDslError::IncompatibleTerm {
                reason:
                    "flexible(...) does not support sas/beta-logistic links; wiggle is only supported for jointly fit standard links"
                        .to_string(),
            });
        }
        return Ok(Some(LinkChoice {
            mode: LinkMode::Flexible,
            link,
            mixture_components: None,
        }));
    }
    if let Some(inner) = t
        .strip_prefix("blended(")
        .and_then(|s| s.strip_suffix(')'))
        .or_else(|| t.strip_prefix("mixture(").and_then(|s| s.strip_suffix(')')))
    {
        if flexible_flag {
            return Err(FormulaDslError::IncompatibleTerm {
                reason:
                    "--flexible-link cannot be combined with --link blended(...)/mixture(...); blended inverse links are not flexible-link mode"
                        .to_string(),
            });
        }
        let components = parse_link_component_list(inner)?;
        return Ok(Some(LinkChoice {
            mode: LinkMode::Strict,
            link: LinkFunction::Logit,
            mixture_components: Some(components),
        }));
    }

    let link = parse_linkname(&t)?;
    if flexible_flag && !linkname_supports_joint_wiggle(link) {
        return Err(FormulaDslError::IncompatibleTerm {
            reason:
                "--flexible-link does not support sas/beta-logistic links; wiggle is only supported for jointly fit standard links"
                    .to_string(),
        });
    }
    Ok(Some(LinkChoice {
        mode: if flexible_flag {
            LinkMode::Flexible
        } else {
            LinkMode::Strict
        },
        link,
        mixture_components: None,
    }))
}

pub fn parse_linkname(v: &str) -> Result<LinkFunction, FormulaDslError> {
    match v.trim() {
        "identity" => Ok(LinkFunction::Identity),
        "log" => Ok(LinkFunction::Log),
        "logit" | "binomial-logit" => Ok(LinkFunction::Logit),
        "probit" | "binomial-probit" => Ok(LinkFunction::Probit),
        "cloglog" | "binomial-cloglog" => Ok(LinkFunction::CLogLog),
        "loglog" => Ok(LinkFunction::LogLog),
        "cauchit" => Ok(LinkFunction::Cauchit),
        "sas" => Ok(LinkFunction::Sas),
        "beta-logistic" => Ok(LinkFunction::BetaLogistic),
        other => Err(FormulaDslError::UnknownIdentifier {
            reason: format!(
                "unsupported link type '{other}'; \
                 use one of identity|log|logit|probit|cloglog|loglog|cauchit|binomial-logit|binomial-probit|binomial-cloglog|sas|beta-logistic|blended(...)/mixture(...) or flexible(...). \
                 Both `--link <type>` (CLI flag) and `link(type=<type>)` (formula term) accept the same set."
            ),
        }),
    }
}

pub fn parse_link_component(v: &str) -> Result<LinkComponent, String> {
    match v.trim() {
        "logit" => Ok(LinkComponent::Logit),
        "probit" => Ok(LinkComponent::Probit),
        "cloglog" => Ok(LinkComponent::CLogLog),
        "loglog" => Ok(LinkComponent::LogLog),
        "cauchit" => Ok(LinkComponent::Cauchit),
        other => Err(FormulaDslError::UnknownIdentifier {
            reason: format!(
                "unsupported blended-link component '{other}'; use probit|logit|cloglog|loglog|cauchit"
            ),
        }
        .into()),
    }
}

pub fn parse_link_component_list(v: &str) -> Result<Vec<LinkComponent>, String> {
    let mut out = Vec::new();
    for part in v.split(',') {
        let trimmed = part.trim();
        if trimmed.is_empty() {
            continue;
        }
        let comp = parse_link_component(trimmed)?;
        if out.contains(&comp) {
            return Err(FormulaDslError::IncompatibleTerm {
                reason: "blended(...) cannot contain duplicate components".to_string(),
            }
            .into());
        }
        out.push(comp);
    }
    if out.len() < 2 {
        return Err(FormulaDslError::InvalidArgument {
            reason: "blended(...) requires at least two components".to_string(),
        }
        .into());
    }
    Ok(out)
}