Skip to main content

gam_terms/inference/
formula_dsl.rs

1use std::collections::{BTreeMap, HashMap};
2
3use pest::Parser;
4use pest::iterators::Pair;
5use pest_derive::Parser;
6
7use crate::smooth::BoundedCoefficientPriorSpec;
8use crate::term_builder::{MARGINAL_SLOPE_Z_ALIAS, marginal_slope_z_alias_is_live};
9use gam_problem::types::{
10    InverseLink, LikelihoodSpec, LinkComponent, LinkFunction, StandardLink, WigglePenaltyConfig,
11};
12
13#[derive(Parser)]
14#[grammar_inline = r#"
15WHITESPACE = _{ " " | "\t" | NEWLINE }
16
17top_function_call = { SOI ~ function_call ~ EOI }
18top_expr = { SOI ~ expr ~ EOI }
19formula = { SOI ~ expr ~ "~" ~ rhs ~ EOI }
20rhs = { term ~ ("+" ~ term)* }
21term = { expr }
22
23expr = { sum }
24sum = { product ~ (add_op ~ product)* }
25add_op = { "+" | "-" }
26product = { interact ~ (mul_op ~ interact)* }
27mul_op = { "*" | "/" }
28interact = { power ~ (interact_op ~ power)* }
29interact_op = { ":" }
30power = { unary ~ (pow_op ~ unary)* }
31pow_op = { "^" }
32unary = { unary_op* ~ primary }
33unary_op = _{ "+" | "-" }
34
35primary = { function_call | list_lit | tuple_lit | ident | number | string_lit | "(" ~ expr ~ ")" }
36list_lit = @{ "[" ~ (!"]" ~ ANY)* ~ "]" }
37tuple_lit = @{ "(" ~ (!("," | ")") ~ ANY)+ ~ "," ~ (!")" ~ ANY)* ~ ")" }
38function_call = { ident ~ "(" ~ arg_list? ~ ")" }
39arg_list = { arg ~ ("," ~ arg)* }
40arg = { named_arg | expr }
41named_arg = { ident ~ "=" ~ expr }
42
43ident = @{ ident_start ~ ident_continue* }
44ident_start = _{ ASCII_ALPHA | "_" }
45ident_continue = _{ ASCII_ALPHANUMERIC | "_" | "." }
46
47number = @{
48    "-"?
49    ~ (ASCII_DIGIT+ ~ ("." ~ ASCII_DIGIT*)? | "." ~ ASCII_DIGIT+)
50    ~ (("e" | "E") ~ ("+" | "-")? ~ ASCII_DIGIT+)?
51}
52
53string_lit = @{ "\"" ~ (!"\"" ~ ANY)* ~ "\"" | "'" ~ (!"'" ~ ANY)* ~ "'" }
54"#]
55struct FormulaParser;
56
57#[derive(Clone, Debug, PartialEq, Eq)]
58pub struct FormulaDslParse {
59    pub response_expr: String,
60    pub rhs_terms: Vec<String>,
61}
62
63#[derive(Clone, Debug, PartialEq, Eq)]
64pub enum CallArgSpec {
65    Positional(String),
66    Named { key: String, value: String },
67}
68
69#[derive(Clone, Debug, PartialEq, Eq)]
70pub struct FunctionCallSpec {
71    pub name: String,
72    pub args: Vec<CallArgSpec>,
73}
74
75/// Typed error surface for the formula DSL parser.
76///
77/// Every variant carries a free-form `reason: String` payload; `Display`
78/// emits exactly that payload, so converting a `FormulaDslError` into
79/// `String` (via the `From` impl below) is byte-equivalent to the pre-
80/// refactor `Err(format!(...))` / `Err("...".to_string())` strings that
81/// the same call sites produced. Public entry points keep their existing
82/// `Result<_, String>` signatures — CLI input handling stays unchanged —
83/// and typed errors flow across the boundary via `From<FormulaDslError>
84/// for String`.
85#[derive(Clone, Debug, PartialEq, Eq)]
86pub enum FormulaDslError {
87    /// Pest grammar failure, unbalanced delimiters, empty terms, or
88    /// missing required parse fragments — i.e. the formula text is
89    /// not a well-formed DSL string.
90    ParseError { reason: String },
91    /// A referenced symbol (link name, blended-link component, term
92    /// function name, top-level RHS identifier) is not part of the
93    /// supported vocabulary.
94    UnknownIdentifier { reason: String },
95    /// A named option's value is unparseable, out of range, or not a
96    /// finite number / valid integer.
97    InvalidArgument { reason: String },
98    /// A combination of terms or options is disallowed (duplicate
99    /// terms, multiple linkwiggle/link/survmodel, mutually exclusive
100    /// option groups in bounded(), wiggle-incompatible links, etc.).
101    IncompatibleTerm { reason: String },
102    /// A required configuration option is missing or empty (e.g.
103    /// `link()` without `type=`, `survmodel()` with no options,
104    /// `bounded()` without a required argument).
105    MalformedConfig { reason: String },
106}
107
108gam_linalg::impl_reason_error_boilerplate! {
109    FormulaDslError {
110        ParseError,
111        UnknownIdentifier,
112        InvalidArgument,
113        IncompatibleTerm,
114        MalformedConfig,
115    }
116}
117
118/// Inbound conversion from `String` is used by `?` cascades inside `parse_formula`
119/// and friends so that internal parser helpers still returning `Result<_, String>`
120/// can flow through without each call site needing an explicit `.map_err(...)`.
121/// We route into `ParseError` because by construction every internal helper that
122/// still produces a raw `String` is itself a parse/term-resolution stage.
123impl From<String> for FormulaDslError {
124    fn from(reason: String) -> Self {
125        FormulaDslError::ParseError { reason }
126    }
127}
128
129pub fn parse_formula_dsl(formula: &str) -> Result<FormulaDslParse, String> {
130    validate_balanced_delimiters(formula, "invalid formula syntax")?;
131    let mut parsed =
132        FormulaParser::parse(Rule::formula, formula).map_err(|e| FormulaDslError::ParseError {
133            reason: format!("invalid formula syntax: {e}"),
134        })?;
135    let formula_pair = parsed.next().ok_or_else(|| FormulaDslError::ParseError {
136        reason: "invalid formula syntax: empty parse".to_string(),
137    })?;
138
139    let mut response_expr: Option<String> = None;
140    let mut rhs_terms: Option<Vec<String>> = None;
141
142    for part in formula_pair.into_inner() {
143        if part.as_rule() == Rule::rhs {
144            rhs_terms = Some(extract_rhs_terms(part)?);
145        } else if part.as_rule() == Rule::expr && response_expr.is_none() {
146            response_expr = Some(part.as_str().trim().to_string());
147        }
148    }
149
150    let response_expr = response_expr.ok_or_else(|| FormulaDslError::ParseError {
151        reason: "invalid formula: missing response expression".to_string(),
152    })?;
153    let rhs_terms = rhs_terms.ok_or_else(|| FormulaDslError::ParseError {
154        reason: "invalid formula: missing RHS terms".to_string(),
155    })?;
156    if rhs_terms.is_empty() {
157        return Err(FormulaDslError::ParseError {
158            reason: "formula has no usable terms".to_string(),
159        }
160        .into());
161    }
162
163    Ok(FormulaDslParse {
164        response_expr,
165        rhs_terms,
166    })
167}
168
169fn delimiter_balance_error(prefix: &str) -> String {
170    format!("{prefix}: unbalanced parentheses or quotes")
171}
172
173// Pest reports malformed delimiters as a generic parse failure. We validate the
174// raw text first so callers get a stable, specific error class for unmatched
175// parentheses/quotes instead of whichever grammar branch happened to fail last.
176fn validate_balanced_delimiters(input: &str, prefix: &str) -> Result<(), String> {
177    let mut stack = Vec::<char>::new();
178    let mut in_single = false;
179    let mut in_double = false;
180
181    for ch in input.chars() {
182        let quoted = in_single || in_double;
183        if ch == '\'' && !in_double {
184            in_single = !in_single;
185        } else if ch == '"' && !in_single {
186            in_double = !in_double;
187        } else if !quoted && matches!(ch, '(' | '[' | '{') {
188            stack.push(ch);
189        } else if !quoted && matches!(ch, ')' | ']' | '}') {
190            let expected = match ch {
191                ')' => '(',
192                ']' => '[',
193                // The enclosing `matches!` guarantees ch is one of ')', ']', '}'.
194                _ => '{',
195            };
196            if stack.pop() != Some(expected) {
197                return Err(FormulaDslError::ParseError {
198                    reason: delimiter_balance_error(prefix),
199                }
200                .into());
201            }
202        }
203    }
204
205    if in_single || in_double || !stack.is_empty() {
206        return Err(FormulaDslError::ParseError {
207            reason: delimiter_balance_error(prefix),
208        }
209        .into());
210    }
211    Ok(())
212}
213
214fn extract_rhs_terms(rhs: Pair<'_, Rule>) -> Result<Vec<String>, String> {
215    let mut out = Vec::new();
216    let mut depth = 0_i32;
217    let mut in_single = false;
218    let mut in_double = false;
219    let mut start = 0_usize;
220    // Last non-whitespace character seen at depth 0 outside quotes. A top-level
221    // `+` only separates terms when it follows a completed operand; when it
222    // follows a binary operator (`:`, `*`, `/`, `^`) or another sign — or opens
223    // the RHS — it is a UNARY sign, not a separator. Splitting on it there would
224    // hand the grammar a truncated fragment like `x:` (from `x:+z`) and surface
225    // a confusing "invalid term syntax in `x:`" instead of the dedicated unary
226    // diagnostic the grammar raises when the whole `x:+z` term reaches it. This
227    // never suppresses a split in a valid formula: no valid Wilkinson-Rogers RHS
228    // places `+` immediately after an operator. `-` is already never treated as
229    // a separator here (it is rejected downstream as a binary term operator), so
230    // the unary `-` path already flows through intact — this restores the same
231    // intact flow for unary `+`.
232    let mut last_significant: Option<char> = None;
233    let text = rhs.as_str();
234    let bytes = text.as_bytes();
235    for (idx, &b) in bytes.iter().enumerate() {
236        let ch = b as char;
237        let quoted = in_single || in_double;
238        if ch == '\'' && !in_double {
239            in_single = !in_single;
240        } else if ch == '"' && !in_single {
241            in_double = !in_double;
242        } else if !quoted && matches!(ch, '(' | '[' | '{') {
243            depth += 1;
244        } else if !quoted && matches!(ch, ')' | ']' | '}') && depth > 0 {
245            depth -= 1;
246        } else if ch == '+'
247            && !quoted
248            && depth == 0
249            && !matches!(
250                last_significant,
251                None | Some(':' | '*' | '/' | '^' | '+' | '-')
252            )
253        {
254            let term = text[start..idx].trim();
255            if term.is_empty() {
256                return Err(FormulaDslError::ParseError {
257                    reason: "formula RHS contains an empty term".to_string(),
258                }
259                .into());
260            }
261            out.push(term.to_string());
262            start = idx + 1;
263        }
264        if !ch.is_ascii_whitespace() {
265            last_significant = Some(ch);
266        }
267    }
268    if in_single || in_double || depth != 0 {
269        return Err(FormulaDslError::ParseError {
270            reason: "formula RHS has unbalanced quotes or parentheses".to_string(),
271        }
272        .into());
273    }
274    let tail = text[start..].trim();
275    if tail.is_empty() {
276        return Err(FormulaDslError::ParseError {
277            reason: "formula RHS contains an empty term".to_string(),
278        }
279        .into());
280    }
281    out.push(tail.to_string());
282    Ok(out)
283}
284
285/// Wilkinson-Rogers operator-family expansion.
286///
287/// A raw RHS term (already split on top-level `+`) may use the documented
288/// formula operators `:`, `*`, `/`, `^`, and parenthesization. This function
289/// expands the term into a normalized list of `WrAtomList`s, where each
290/// `WrAtomList` is one resulting model term — either a single atom (linear
291/// main effect / function-call term) or multiple atoms (interaction).
292///
293/// Wilkinson-Rogers semantics implemented here:
294/// * `a` produces `{a}` — one main effect.
295/// * `a:b` produces `{a:b}` — one interaction.
296/// * `a*b` produces `{a, b, a:b}` — crossing (expanded).
297/// * `a/b` produces `{a, a:b}` — nesting. Nesting is hierarchical: each `/`
298///   crosses its right operand with the product of ALL variables on its left,
299///   so `a/b/c` produces `{a, a:b, a:b:c}` (never a spurious `a:c`) and
300///   `a*b/c` produces `{a, b, a:b, a:b:c}`, matching R/mgcv `terms()`.
301/// * `(a + b + ... )^n` produces every non-empty subset of `{a, b, ...}` of
302///   size at most `n`, each subset being one interaction.
303/// * `+` unions two term sets; `-` is rejected.
304///
305/// Atoms inside the AST may be bare identifiers or function calls. Function
306/// calls are opaque — they pass through as-is and cannot participate in an
307/// `:` interaction with other atoms (smooths/factors require dedicated
308/// constructors like `te()`).
309type WrAtomList = Vec<String>;
310
311fn expand_wr_term(raw: &str) -> Result<Vec<WrAtomList>, String> {
312    let mut parsed = FormulaParser::parse(Rule::top_expr, raw).map_err(|e| {
313        FormulaDslError::ParseError {
314            reason: format!("invalid term syntax in `{raw}`: {e}"),
315        }
316        .to_string()
317    })?;
318    let top = parsed.next().ok_or_else(|| {
319        FormulaDslError::ParseError {
320            reason: format!("invalid term syntax in `{raw}`: empty parse"),
321        }
322        .to_string()
323    })?;
324    let expr = top
325        .into_inner()
326        .find(|p| p.as_rule() == Rule::expr)
327        .ok_or_else(|| {
328            FormulaDslError::ParseError {
329                reason: format!("invalid term syntax in `{raw}`: missing expr"),
330            }
331            .to_string()
332        })?;
333    let interactions = expand_expr(expr, raw)?;
334    let normalized: Vec<WrAtomList> = interactions
335        .into_iter()
336        .map(normalize_interaction)
337        .collect();
338    // Deduplicate while preserving order: a*b yields {a, b, a:b}; a*b + a
339    // would otherwise list `a` twice.
340    let mut seen = std::collections::BTreeSet::<Vec<String>>::new();
341    let mut out = Vec::<WrAtomList>::new();
342    for term in normalized {
343        let key = term.clone();
344        if seen.insert(key) {
345            out.push(term);
346        }
347    }
348    Ok(out)
349}
350
351fn normalize_interaction(mut atoms: WrAtomList) -> WrAtomList {
352    atoms.sort();
353    atoms.dedup();
354    atoms
355}
356
357fn expand_expr(pair: Pair<'_, Rule>, raw: &str) -> Result<Vec<WrAtomList>, String> {
358    match pair.as_rule() {
359        Rule::expr => {
360            let inner = pair.into_inner().next().ok_or_else(|| {
361                FormulaDslError::ParseError {
362                    reason: format!("invalid term syntax in `{raw}`: empty expr"),
363                }
364                .to_string()
365            })?;
366            expand_expr(inner, raw)
367        }
368        Rule::sum => {
369            let mut iter = pair.into_inner();
370            let first = iter.next().ok_or_else(|| {
371                FormulaDslError::ParseError {
372                    reason: format!("invalid term syntax in `{raw}`: empty sum"),
373                }
374                .to_string()
375            })?;
376            let mut acc = expand_expr(first, raw)?;
377            while let Some(op) = iter.next() {
378                if op.as_rule() != Rule::add_op {
379                    return Err(FormulaDslError::ParseError {
380                        reason: format!(
381                            "invalid term syntax in `{raw}`: expected add operator, got `{:?}`",
382                            op.as_rule()
383                        ),
384                    }
385                    .into());
386                }
387                let op_str = op.as_str().trim();
388                let operand = iter.next().ok_or_else(|| {
389                    FormulaDslError::ParseError {
390                        reason: format!("invalid term syntax in `{raw}`: dangling `{op_str}`"),
391                    }
392                    .to_string()
393                })?;
394                if op_str == "-" {
395                    return Err(FormulaDslError::IncompatibleTerm {
396                        reason: format!(
397                            "binary `-` is not supported inside a formula term in `{raw}` \
398                             (use multiple `+` terms or drop the unwanted predictor explicitly)"
399                        ),
400                    }
401                    .into());
402                }
403                let mut rhs = expand_expr(operand, raw)?;
404                acc.append(&mut rhs);
405            }
406            Ok(acc)
407        }
408        Rule::product => {
409            let mut iter = pair.into_inner();
410            let first = iter.next().ok_or_else(|| {
411                FormulaDslError::ParseError {
412                    reason: format!("invalid term syntax in `{raw}`: empty product"),
413                }
414                .to_string()
415            })?;
416            let mut acc = expand_expr(first, raw)?;
417            while let Some(op) = iter.next() {
418                if op.as_rule() != Rule::mul_op {
419                    return Err(FormulaDslError::ParseError {
420                        reason: format!(
421                            "invalid term syntax in `{raw}`: expected `*` or `/`, got `{:?}`",
422                            op.as_rule()
423                        ),
424                    }
425                    .into());
426                }
427                let op_str = op.as_str().trim();
428                let operand = iter.next().ok_or_else(|| {
429                    FormulaDslError::ParseError {
430                        reason: format!("invalid term syntax in `{raw}`: dangling `{op_str}`"),
431                    }
432                    .to_string()
433                })?;
434                let rhs = expand_expr(operand, raw)?;
435                acc = match op_str {
436                    "*" => wr_cross(acc, rhs),
437                    "/" => wr_nest(acc, rhs),
438                    other => {
439                        return Err(FormulaDslError::ParseError {
440                            reason: format!(
441                                "invalid term syntax in `{raw}`: unrecognized mul operator `{other}`"
442                            ),
443                        }
444                        .into());
445                    }
446                };
447            }
448            Ok(acc)
449        }
450        Rule::interact => {
451            let mut iter = pair.into_inner();
452            let first = iter.next().ok_or_else(|| {
453                FormulaDslError::ParseError {
454                    reason: format!("invalid term syntax in `{raw}`: empty interact"),
455                }
456                .to_string()
457            })?;
458            let mut acc = expand_expr(first, raw)?;
459            while let Some(op) = iter.next() {
460                if op.as_rule() != Rule::interact_op {
461                    return Err(FormulaDslError::ParseError {
462                        reason: format!(
463                            "invalid term syntax in `{raw}`: expected `:`, got `{:?}`",
464                            op.as_rule()
465                        ),
466                    }
467                    .into());
468                }
469                let operand = iter.next().ok_or_else(|| {
470                    FormulaDslError::ParseError {
471                        reason: format!("invalid term syntax in `{raw}`: dangling `:`"),
472                    }
473                    .to_string()
474                })?;
475                let rhs = expand_expr(operand, raw)?;
476                acc = wr_interact(acc, rhs, raw)?;
477            }
478            Ok(acc)
479        }
480        Rule::power => {
481            let mut iter = pair.into_inner();
482            let first = iter.next().ok_or_else(|| {
483                FormulaDslError::ParseError {
484                    reason: format!("invalid term syntax in `{raw}`: empty power"),
485                }
486                .to_string()
487            })?;
488            let base = expand_expr(first, raw)?;
489            let Some(op) = iter.next() else {
490                return Ok(base);
491            };
492            if op.as_rule() != Rule::pow_op {
493                return Err(FormulaDslError::ParseError {
494                    reason: format!(
495                        "invalid term syntax in `{raw}`: expected `^`, got `{:?}`",
496                        op.as_rule()
497                    ),
498                }
499                .into());
500            }
501            let exponent_pair = iter.next().ok_or_else(|| {
502                FormulaDslError::ParseError {
503                    reason: format!("invalid term syntax in `{raw}`: dangling `^`"),
504                }
505                .to_string()
506            })?;
507            let exp_text = exponent_pair.as_str().trim();
508            let n: usize = exp_text.parse().map_err(|_| {
509                FormulaDslError::ParseError {
510                    reason: format!(
511                        "invalid term syntax in `{raw}`: `^` exponent must be a positive integer, got `{exp_text}`"
512                    ),
513                }
514                .to_string()
515            })?;
516            if n == 0 {
517                return Err(FormulaDslError::ParseError {
518                    reason: format!(
519                        "invalid term syntax in `{raw}`: `^0` is not a meaningful formula expansion"
520                    ),
521                }
522                .into());
523            }
524            if let Some(extra_op) = iter.next() {
525                let extra = extra_op.as_str().trim();
526                return Err(FormulaDslError::ParseError {
527                    reason: format!(
528                        "invalid term syntax in `{raw}`: chained `^` operators are not supported; \
529                         use one positive integer exponent, got another `{extra}`"
530                    ),
531                }
532                .into());
533            }
534            Ok(wr_power(base, n))
535        }
536        Rule::unary => {
537            let unary_start = pair.as_span().start();
538            for inner in pair.into_inner() {
539                if inner.as_rule() == Rule::primary {
540                    let prefix_len = inner.as_span().start().saturating_sub(unary_start);
541                    if prefix_len > 0 {
542                        let prefix = &raw[unary_start..inner.as_span().start()];
543                        if prefix.chars().any(|ch| matches!(ch, '+' | '-')) {
544                            return Err(FormulaDslError::IncompatibleTerm {
545                                reason: format!(
546                                    "unary `+`/`-` is not supported inside a formula term in `{raw}`"
547                                ),
548                            }
549                            .into());
550                        }
551                    }
552                    return expand_expr(inner, raw);
553                }
554            }
555            Err(FormulaDslError::ParseError {
556                reason: format!("invalid term syntax in `{raw}`: empty unary"),
557            }
558            .into())
559        }
560        Rule::primary => {
561            let span = pair.as_str().trim().to_string();
562            let inner = pair.into_inner().next();
563            match inner {
564                Some(child) if child.as_rule() == Rule::expr => expand_expr(child, raw),
565                Some(child) => Ok(vec![vec![child.as_str().trim().to_string()]]),
566                None => Ok(vec![vec![span]]),
567            }
568        }
569        _ => Err(FormulaDslError::ParseError {
570            reason: format!(
571                "invalid term syntax in `{raw}`: unexpected node `{:?}`",
572                pair.as_rule()
573            ),
574        }
575        .into()),
576    }
577}
578
579fn wr_cross(left: Vec<WrAtomList>, right: Vec<WrAtomList>) -> Vec<WrAtomList> {
580    // a*b = a + b + a:b
581    let mut out = Vec::with_capacity(left.len() + right.len() + left.len() * right.len());
582    out.extend(left.iter().cloned());
583    out.extend(right.iter().cloned());
584    for l in &left {
585        for r in &right {
586            let mut merged: WrAtomList = l.iter().cloned().chain(r.iter().cloned()).collect();
587            merged.sort();
588            merged.dedup();
589            out.push(merged);
590        }
591    }
592    out
593}
594
595fn wr_nest(left: Vec<WrAtomList>, right: Vec<WrAtomList>) -> Vec<WrAtomList> {
596    // Wilkinson-Rogers nesting `A/B`: keep every term of `A`, then nest `B`
597    // within the *whole* left-hand group by crossing each right term with the
598    // product of ALL variables appearing anywhere on the left (not with each
599    // left term separately). This makes chained nesting hierarchical —
600    // `a/b/c` = {a, a:b, a:b:c}, never a spurious `a:c` — and matches R/mgcv
601    // `terms()` exactly, including `a*b/c` = {a, b, a:b, a:b:c} and
602    // `(a+b)/c` = {a, b, a:b:c}. See `parse_formula_chained_wr_nesting_is_hierarchical`.
603    let mut left_atoms: WrAtomList = left.iter().flatten().cloned().collect();
604    left_atoms.sort();
605    left_atoms.dedup();
606
607    let mut out = Vec::with_capacity(left.len() + right.len());
608    out.extend(left.iter().cloned());
609    for r in &right {
610        let mut merged: WrAtomList = left_atoms
611            .iter()
612            .cloned()
613            .chain(r.iter().cloned())
614            .collect();
615        merged.sort();
616        merged.dedup();
617        out.push(merged);
618    }
619    out
620}
621
622fn wr_interact(
623    left: Vec<WrAtomList>,
624    right: Vec<WrAtomList>,
625    raw: &str,
626) -> Result<Vec<WrAtomList>, String> {
627    // a:b — Cartesian merge of every (l, r) pair.
628    let mut out = Vec::with_capacity(left.len() * right.len());
629    for l in &left {
630        for r in &right {
631            let combined: WrAtomList = l.iter().cloned().chain(r.iter().cloned()).collect();
632            // Reject interactions whose atoms contain function-call syntax.
633            // These would build design columns that are not simple products
634            // (factors, smooths) and need a dedicated constructor.
635            for atom in &combined {
636                if atom.contains('(') {
637                    return Err(FormulaDslError::IncompatibleTerm {
638                        reason: format!(
639                            "interaction operator `:` with function-call atom is not supported in `{raw}`. \
640                             Use te(...) for smooth interactions or group()/factor() with a separate \
641                             interaction strategy for categorical effects."
642                        ),
643                    }
644                    .into());
645                }
646            }
647            let mut merged = combined;
648            merged.sort();
649            merged.dedup();
650            out.push(merged);
651        }
652    }
653    Ok(out)
654}
655
656fn wr_power(base: Vec<WrAtomList>, n: usize) -> Vec<WrAtomList> {
657    // (e)^n produces every non-empty subset of `base` of size ≤ n,
658    // each subset being one interaction. The standard WR semantics
659    // treat the base as a *set* of atoms — interactions inside `base`
660    // are kept as-is and not re-crossed.
661    if base.is_empty() {
662        return Vec::new();
663    }
664    let m = base.len();
665    let mut out = Vec::<WrAtomList>::new();
666    let max_size = n.min(m);
667    for size in 1..=max_size {
668        // Enumerate combinations of `size` elements out of m.
669        let mut indices: Vec<usize> = (0..size).collect();
670        loop {
671            let mut merged = WrAtomList::new();
672            for &i in &indices {
673                merged.extend(base[i].iter().cloned());
674            }
675            merged.sort();
676            merged.dedup();
677            out.push(merged);
678            // Next combination
679            let mut k = size;
680            while k > 0 {
681                k -= 1;
682                if indices[k] != k + m - size {
683                    indices[k] += 1;
684                    for j in (k + 1)..size {
685                        indices[j] = indices[j - 1] + 1;
686                    }
687                    break;
688                }
689                if k == 0 {
690                    k = usize::MAX;
691                    break;
692                }
693            }
694            if k == usize::MAX {
695                break;
696            }
697        }
698    }
699    out
700}
701
702fn is_exact_ident(raw: &str) -> bool {
703    let mut chars = raw.chars();
704    let Some(first) = chars.next() else {
705        return false;
706    };
707    if !first.is_ascii_alphabetic() && first != '_' {
708        return false;
709    }
710    chars.all(|ch| ch.is_ascii_alphanumeric() || ch == '_' || ch == '.')
711}
712
713pub fn parse_function_call(input: &str) -> Result<FunctionCallSpec, String> {
714    validate_balanced_delimiters(input, "invalid function call syntax")?;
715    let mut parsed = FormulaParser::parse(Rule::top_function_call, input).map_err(|e| {
716        FormulaDslError::ParseError {
717            reason: format!("invalid function call syntax: {e}"),
718        }
719    })?;
720    let top = parsed.next().ok_or_else(|| FormulaDslError::ParseError {
721        reason: "invalid function call syntax: empty parse".to_string(),
722    })?;
723    let call = top
724        .into_inner()
725        .find(|p| p.as_rule() == Rule::function_call)
726        .ok_or_else(|| FormulaDslError::ParseError {
727            reason: "invalid function call syntax: missing call".to_string(),
728        })?;
729    parse_call_pair(call)
730}
731
732fn parse_call_pair(call: Pair<'_, Rule>) -> Result<FunctionCallSpec, String> {
733    let mut name: Option<String> = None;
734    let mut args = Vec::<CallArgSpec>::new();
735    for part in call.into_inner() {
736        if part.as_rule() == Rule::ident {
737            if name.is_none() {
738                name = Some(part.as_str().trim().to_string());
739            }
740        } else if part.as_rule() == Rule::arg_list {
741            for a in part.into_inner() {
742                if a.as_rule() != Rule::arg {
743                    continue;
744                }
745                let mut a_inner = a.into_inner();
746                let Some(first) = a_inner.next() else {
747                    continue;
748                };
749                if first.as_rule() == Rule::named_arg {
750                    let mut ni = first.into_inner();
751                    let key = ni
752                        .next()
753                        .ok_or_else(|| FormulaDslError::ParseError {
754                            reason: "invalid named argument key".to_string(),
755                        })?
756                        .as_str()
757                        .trim()
758                        .to_ascii_lowercase();
759                    let value = ni
760                        .next()
761                        .ok_or_else(|| FormulaDslError::ParseError {
762                            reason: "invalid named argument value".to_string(),
763                        })?
764                        .as_str()
765                        .trim()
766                        .to_string();
767                    args.push(CallArgSpec::Named { key, value });
768                } else if first.as_rule() == Rule::expr {
769                    args.push(CallArgSpec::Positional(first.as_str().trim().to_string()));
770                }
771            }
772        }
773    }
774    let name = name.ok_or_else(|| FormulaDslError::ParseError {
775        reason: "invalid function call: missing name".to_string(),
776    })?;
777    Ok(FunctionCallSpec { name, args })
778}
779
780#[cfg(test)]
781mod tests {
782    use super::{
783        CallArgSpec, ParsedTerm, parse_formula, parse_formula_dsl, parse_function_call,
784        parse_linkwiggle_formulaspec, parsed_term_column_names, parsed_terms_reference_column,
785        validate_marginal_slope_z_alias_exclusion, validate_marginal_slope_z_column_exclusion,
786    };
787    use std::collections::{BTreeMap, BTreeSet, HashMap};
788
789    #[test]
790    fn parsed_term_column_names_includes_by_smooth_grouping_variable() {
791        // The model's input contract must include a smooth's `by=` column, not
792        // just its positional variable. `s(x, by=g)` consumes both `x` and `g`
793        // (`term_builder` reads `options["by"]`); dropping `g` from the
794        // required/consumable set would either omit a genuine predictor at fit
795        // (CLI projected load) or — worse — make predict silently project the
796        // `g` column away (#840 regression). A bare main effect and an
797        // interaction round out the variant coverage.
798        let parsed =
799            parse_formula("y ~ s(x, by=g) + z + a:b").expect("formula with a by= smooth parses");
800        let mut cols = BTreeSet::<String>::new();
801        parsed_term_column_names(&parsed.terms, &mut cols);
802        for expected in ["x", "g", "z", "a", "b"] {
803            assert!(
804                cols.contains(expected),
805                "parsed_term_column_names dropped '{expected}'; got {cols:?}"
806            );
807        }
808        // The response is never a *term* column (it is handled separately).
809        assert!(
810            !cols.contains("y"),
811            "response leaked into term columns: {cols:?}"
812        );
813    }
814
815    #[test]
816    fn linkwiggle_parser_does_not_bake_in_cubic_only_restriction() {
817        // Regression for #384: the cubic-only constraint belongs to the
818        // score-warp / link-deviation `DeviationRuntime`, NOT to this shared
819        // parser. `parse_linkwiggle_formulaspec` also feeds `timewiggle` and
820        // the location-scale survival path, whose general monotone I-spline
821        // value basis honors any `degree >= 2`. So the parser must accept
822        // non-cubic degrees and carry them through verbatim; the cubic gate is
823        // applied downstream only where the cubic-only runtime is built. A
824        // prior fix wrongly forced `degree == 3` here, breaking timewiggle and
825        // location-scale callers — this pins that the parser stays general.
826        for deg in [2usize, 4, 5, 10] {
827            let mut options = BTreeMap::new();
828            options.insert("degree".to_string(), deg.to_string());
829            options.insert("internal_knots".to_string(), "3".to_string());
830            let raw = format!("timewiggle(degree={deg}, internal_knots=3)");
831            let spec = parse_linkwiggle_formulaspec(&options, &raw)
832                .expect("non-cubic wiggle degree must parse at the shared layer");
833            assert_eq!(
834                spec.degree, deg,
835                "parser must carry the requested degree through verbatim"
836            );
837        }
838
839        // The only universal lower bound the shared parser enforces is that a
840        // polynomial degree is positive.
841        let mut zero = BTreeMap::new();
842        zero.insert("degree".to_string(), "0".to_string());
843        zero.insert("internal_knots".to_string(), "3".to_string());
844        let err = parse_linkwiggle_formulaspec(&zero, "linkwiggle(degree=0, internal_knots=3)")
845            .expect_err("degree=0 must be rejected");
846        assert!(
847            err.contains("degree >= 1"),
848            "error should state the positive-degree lower bound, got: {err}"
849        );
850    }
851
852    #[test]
853    fn parses_nested_formula_terms() {
854        let parsed =
855            parse_formula_dsl("log(y) ~ x1 + s(log(x2 + 1), bs=\"tps\", k=10) + te(x3, x4)")
856                .expect("parse");
857        assert_eq!(parsed.response_expr, "log(y)");
858        assert_eq!(parsed.rhs_terms.len(), 3);
859        assert_eq!(parsed.rhs_terms[0], "x1");
860        assert_eq!(parsed.rhs_terms[1], "s(log(x2 + 1), bs=\"tps\", k=10)");
861        assert_eq!(parsed.rhs_terms[2], "te(x3, x4)");
862    }
863
864    #[test]
865    fn parses_cyclic_formula_aliases() {
866        let parsed = parse_formula("y ~ cyclic(theta, period_start=0, period_end=6.283)")
867            .expect("parse cyclic formula");
868        match &parsed.terms[0] {
869            super::ParsedTerm::Smooth { vars, options, .. } => {
870                assert_eq!(vars, &vec!["theta".to_string()]);
871                assert_eq!(options.get("type").map(String::as_str), Some("cyclic"));
872                assert_eq!(options.get("period_start").map(String::as_str), Some("0"));
873            }
874            other => panic!("expected cyclic smooth term, got {other:?}"),
875        }
876    }
877
878    #[test]
879    fn sphere_aliases_all_dispatch_to_intrinsic_s2_basis() {
880        // Regression for #383: `s2(lat, lon)` must build the same intrinsic
881        // S² (sphere) basis as `sphere()`/`sos()`/`spherical()`. Previously the
882        // `s2` arm returned a Smooth without `type=sphere`, so it silently fell
883        // back to a generic Euclidean 2-D smooth and diverged in the
884        // spatial-kappa optimizer. All four aliases must be byte-for-byte
885        // equivalent in their dispatch (vars + `type=sphere`).
886        for alias in ["sphere", "sos", "spherical", "s2"] {
887            let parsed = parse_formula(&format!("y ~ {alias}(lat, lon)"))
888                .unwrap_or_else(|e| panic!("parse {alias}: {e}"));
889            match &parsed.terms[0] {
890                super::ParsedTerm::Smooth { vars, options, .. } => {
891                    assert_eq!(
892                        vars,
893                        &vec!["lat".to_string(), "lon".to_string()],
894                        "{alias} should keep (lat, lon) as its variables"
895                    );
896                    assert_eq!(
897                        options.get("type").map(String::as_str),
898                        Some("sphere"),
899                        "{alias} must dispatch to the intrinsic sphere basis (type=sphere)"
900                    );
901                }
902                other => panic!("expected sphere smooth term for {alias}, got {other:?}"),
903            }
904        }
905    }
906
907    #[test]
908    fn parses_function_callwithnamed_and_positional_args() {
909        let call = parse_function_call("s(log(x + 1), type=\"duchon\", centers=12)").expect("call");
910        assert_eq!(call.name, "s");
911        assert_eq!(call.args.len(), 3);
912        assert_eq!(
913            call.args[0],
914            CallArgSpec::Positional("log(x + 1)".to_string())
915        );
916        assert_eq!(
917            call.args[1],
918            CallArgSpec::Named {
919                key: "type".to_string(),
920                value: "\"duchon\"".to_string()
921            }
922        );
923    }
924
925    #[test]
926    fn parses_tensor_boundary_list_options() {
927        let call = parse_function_call(
928            "te(day_of_week, hour, boundary=['periodic', 'periodic'], period=[7, 24])",
929        )
930        .expect("call");
931        assert_eq!(call.name, "te");
932        assert_eq!(call.args.len(), 4);
933        assert_eq!(
934            call.args[2],
935            CallArgSpec::Named {
936                key: "boundary".to_string(),
937                value: "['periodic', 'periodic']".to_string(),
938            }
939        );
940    }
941
942    #[test]
943    fn parse_formula_dsl_reports_unbalanced_parentheses() {
944        let err = parse_formula_dsl("y ~ s(x, k=10").expect_err("expected parse failure");
945        assert!(err.contains("unbalanced parentheses"));
946    }
947
948    #[test]
949    fn parse_function_call_reports_unbalanced_parentheses() {
950        let err = parse_function_call("s(x, k=10").expect_err("expected parse failure");
951        assert!(err.contains("unbalanced parentheses"));
952    }
953
954    #[test]
955    fn parse_formula_accepts_tuple_smooth_options() {
956        let parsed = parse_formula("z ~ te(x, y, k=(20, 20))")
957            .expect("tuple-valued smooth option should parse");
958        assert_eq!(parsed.terms.len(), 1);
959
960        let dsl = parse_formula_dsl("z ~ te(x, y, k=(20, 20))")
961            .expect("tuple-valued smooth option should parse in the DSL layer");
962        assert_eq!(dsl.rhs_terms, vec!["te(x, y, k=(20, 20))"]);
963
964        let call = parse_function_call("te(x, y, k=(20, 20))")
965            .expect("tuple-valued smooth option should parse as a function call");
966        assert_eq!(
967            call.args[2],
968            CallArgSpec::Named {
969                key: "k".to_string(),
970                value: "(20, 20)".to_string(),
971            }
972        );
973    }
974
975    #[test]
976    fn parse_formula_rejects_unsupported_top_level_rhs_expressions() {
977        // Binary `-`, unary `-`, bare parens, and `-1` intercept removal are
978        // not supported as top-level RHS expressions; they must surface
979        // through the bare-identifier check in `parse_term` with the same
980        // diagnostic. WR operators `:`, `*`, `/`, and `^` are intentionally
981        // supported by `expand_wr_term` and are exercised by
982        // `parse_formula_supports_wr_slash_nesting` /
983        // `parse_formula_supports_wr_star_crossing`; they must NOT appear in
984        // this list.
985        for formula in ["y ~ x - z", "y ~ -x", "y ~ (x)", "y ~ x - 1"] {
986            let err = parse_formula(formula).expect_err("expected formula parse failure");
987            assert!(err.to_string().contains("unsupported top-level RHS term"));
988        }
989    }
990
991    /// Wilkinson-Rogers `a/b` (nesting) is documented and implemented by
992    /// `expand_wr_term` to yield `{a, a:b}`. This test pins that contract end
993    /// to end so a regression in either the grammar's `mul_op = { "*" | "/" }`
994    /// branch or in `expand_expr`'s `Rule::product` handling surfaces here
995    /// instead of slipping into a silent re-interpretation as a rejection.
996    #[test]
997    fn parse_formula_supports_wr_slash_nesting() {
998        let parsed = parse_formula("y ~ x / z").expect("`/` is supported as WR nesting");
999        assert_eq!(parsed.response, "y");
1000        assert_eq!(parsed.terms.len(), 2);
1001        let names: Vec<String> = parsed
1002            .terms
1003            .iter()
1004            .map(|t| match t {
1005                ParsedTerm::Linear { name, .. } => format!("Linear({name})"),
1006                ParsedTerm::Interaction { vars, .. } => {
1007                    format!("Interaction({})", vars.join(":"))
1008                }
1009                other => format!("Other({other:?})"),
1010            })
1011            .collect();
1012        assert_eq!(
1013            names,
1014            vec!["Linear(x)".to_string(), "Interaction(x:z)".to_string()]
1015        );
1016    }
1017
1018    /// Collect the WR term labels of a parsed formula in `:`-joined form,
1019    /// matching R/mgcv `attr(terms(...), "term.labels")` so nesting/crossing
1020    /// expansions can be pinned against ground truth.
1021    fn wr_term_labels(formula: &str) -> Vec<String> {
1022        let parsed = parse_formula(formula).unwrap_or_else(|e| panic!("parse {formula}: {e}"));
1023        parsed
1024            .terms
1025            .iter()
1026            .map(|t| match t {
1027                ParsedTerm::Linear { name, .. } => name.clone(),
1028                ParsedTerm::Interaction { vars, .. } => vars.join(":"),
1029                other => format!("Other({other:?})"),
1030            })
1031            .collect()
1032    }
1033
1034    /// Regression for #2290: chained Wilkinson-Rogers nesting `a/b/c` must be
1035    /// hierarchical — the deepest interaction crosses `c` with the FULL `a:b`,
1036    /// giving `{a, a:b, a:b:c}` with NO spurious `a:c`. Before the fix,
1037    /// `Rule::product`'s left-associative fold distributed `c` over every
1038    /// accumulated left term (`{a, a:b}`), minting `a:c`. Each expansion is
1039    /// pinned against R/mgcv `terms()` ground truth.
1040    #[test]
1041    fn parse_formula_chained_wr_nesting_is_hierarchical() {
1042        // R: terms(y ~ a/b/c) -> a  a:b  a:b:c
1043        assert_eq!(
1044            wr_term_labels("y ~ a/b/c"),
1045            vec!["a".to_string(), "a:b".to_string(), "a:b:c".to_string()],
1046            "a/b/c must nest hierarchically with no spurious a:c"
1047        );
1048        // R: terms(y ~ a*b/c) -> a  b  a:b  a:b:c
1049        assert_eq!(
1050            wr_term_labels("y ~ a*b/c"),
1051            vec![
1052                "a".to_string(),
1053                "b".to_string(),
1054                "a:b".to_string(),
1055                "a:b:c".to_string()
1056            ],
1057            "a*b/c nests c within the whole a*b group"
1058        );
1059        // R: terms(y ~ x/z) -> x  x:z  (single-level nesting unchanged)
1060        assert_eq!(
1061            wr_term_labels("y ~ x/z"),
1062            vec!["x".to_string(), "x:z".to_string()],
1063        );
1064        // R: terms(y ~ a/b/c/d) -> a  a:b  a:b:c  a:b:c:d
1065        assert_eq!(
1066            wr_term_labels("y ~ a/b/c/d"),
1067            vec![
1068                "a".to_string(),
1069                "a:b".to_string(),
1070                "a:b:c".to_string(),
1071                "a:b:c:d".to_string()
1072            ],
1073        );
1074    }
1075
1076    /// Wilkinson-Rogers `a*b` (crossing) is documented and implemented by
1077    /// `expand_wr_term` to yield `{a, b, a:b}`. Pin the contract so a future
1078    /// refactor of `mul_op` / `Rule::product` handling cannot silently change
1079    /// the set of model terms.
1080    #[test]
1081    fn parse_formula_supports_wr_star_crossing() {
1082        let parsed = parse_formula("y ~ x * z").expect("`*` is supported as WR crossing");
1083        assert_eq!(parsed.response, "y");
1084        assert_eq!(parsed.terms.len(), 3);
1085        let names: Vec<String> = parsed
1086            .terms
1087            .iter()
1088            .map(|t| match t {
1089                ParsedTerm::Linear { name, .. } => format!("Linear({name})"),
1090                ParsedTerm::Interaction { vars, .. } => {
1091                    format!("Interaction({})", vars.join(":"))
1092                }
1093                other => format!("Other({other:?})"),
1094            })
1095            .collect();
1096        assert_eq!(
1097            names,
1098            vec![
1099                "Linear(x)".to_string(),
1100                "Linear(z)".to_string(),
1101                "Interaction(x:z)".to_string(),
1102            ]
1103        );
1104    }
1105
1106    #[test]
1107    fn parse_formula_rejects_unary_signs_inside_wr_expansion() {
1108        for formula in ["y ~ x:-z", "y ~ a*-b", "y ~ x/-z", "y ~ x:+z"] {
1109            let err = parse_formula(formula)
1110                .expect_err("WR expansion must not silently drop unary signs");
1111            let msg = err.to_string();
1112            assert!(
1113                msg.contains("unary `+`/`-` is not supported"),
1114                "unexpected error for {formula}: {msg}"
1115            );
1116        }
1117    }
1118
1119    #[test]
1120    fn parse_formula_supports_wr_power_crossing() {
1121        let parsed = parse_formula("y ~ (x + z)^2").expect("`^` is supported as WR power");
1122        assert_eq!(parsed.response, "y");
1123        assert_eq!(parsed.terms.len(), 3);
1124        let names: Vec<String> = parsed
1125            .terms
1126            .iter()
1127            .map(|t| match t {
1128                ParsedTerm::Linear { name, .. } => format!("Linear({name})"),
1129                ParsedTerm::Interaction { vars, .. } => {
1130                    format!("Interaction({})", vars.join(":"))
1131                }
1132                other => format!("Other({other:?})"),
1133            })
1134            .collect();
1135        assert_eq!(
1136            names,
1137            vec![
1138                "Linear(x)".to_string(),
1139                "Linear(z)".to_string(),
1140                "Interaction(x:z)".to_string(),
1141            ]
1142        );
1143    }
1144
1145    #[test]
1146    fn parse_formula_rejects_chained_wr_power() {
1147        let err = parse_formula("y ~ (x + z)^2^3")
1148            .expect_err("chained WR powers must not silently drop later exponents");
1149        let msg = err.to_string();
1150        assert!(
1151            msg.contains("chained `^` operators are not supported"),
1152            "error should explain that chained powers are rejected, got: {msg}"
1153        );
1154    }
1155
1156    #[test]
1157    fn parsed_terms_reference_column_sees_the_by_smooth_variable() {
1158        // Regression for #807: the by= grouping variable lives in
1159        // options["by"], not the smooth's positional vars. The reference
1160        // predicate must still recognise it, both so the CLI loads the column
1161        // and so the marginal-slope z-column exclusion check (which reuses this
1162        // predicate) cannot be fooled into aliasing a reserved z onto a by=.
1163        let parsed = parse_formula("y ~ s(x, by=g)").expect("parse by-smooth");
1164        assert!(
1165            parsed_terms_reference_column(&parsed.terms, "g"),
1166            "s(x, by=g) references column g via options[\"by\"]"
1167        );
1168        assert!(parsed_terms_reference_column(&parsed.terms, "x"));
1169        assert!(!parsed_terms_reference_column(&parsed.terms, "absent"));
1170    }
1171
1172    #[test]
1173    fn marginal_slope_z_column_validator_detects_linear_and_smooth_reuse() {
1174        let main = parse_formula("y ~ x + z").expect("parse main");
1175        let logslope = parse_formula("y ~ s(z, type=duchon, centers=6)").expect("parse logslope");
1176
1177        assert!(parsed_terms_reference_column(&main.terms, "z"));
1178        assert!(parsed_terms_reference_column(&logslope.terms, "z"));
1179
1180        let err = validate_marginal_slope_z_column_exclusion(
1181            &main,
1182            &parse_formula("y ~ 1").expect("parse clean logslope"),
1183            "z",
1184            "bernoulli marginal-slope",
1185            "--logslope-formula",
1186        )
1187        .expect_err("main formula should be rejected");
1188        assert!(err.contains("cannot also appear in the main formula"));
1189
1190        let err = validate_marginal_slope_z_column_exclusion(
1191            &parse_formula("y ~ x").expect("parse clean main"),
1192            &logslope,
1193            "z",
1194            "bernoulli marginal-slope",
1195            "--logslope-formula",
1196        )
1197        .expect_err("logslope formula should be rejected");
1198        assert!(err.contains("cannot also appear in --logslope-formula"));
1199    }
1200
1201    /// Column map helper for the alias tests: names -> positional indices.
1202    fn frame(columns: &[&str]) -> HashMap<String, usize> {
1203        columns
1204            .iter()
1205            .enumerate()
1206            .map(|(i, c)| ((*c).to_string(), i))
1207            .collect()
1208    }
1209
1210    /// The case the literal-name validator is structurally blind to.
1211    ///
1212    /// Every pre-existing test above uses `z_column = "z"`, where the reserved
1213    /// name and the canonical alias are the same string — so none of them can
1214    /// distinguish a validator that resolves the alias from one that only
1215    /// compares text. With `z_column != "z"` they come apart, and `y ~ x + z`
1216    /// silently put the score into the marginal design (gam#2432).
1217    #[test]
1218    fn marginal_slope_alias_validator_rejects_bare_z_when_z_column_is_named_otherwise() {
1219        let cols = frame(&["y", "x", "pgs_ctn_z"]);
1220        let main = parse_formula("y ~ x + z").expect("parse main");
1221
1222        // The literal-name validator accepts it: the term is named `z`, the
1223        // reserved column is `pgs_ctn_z`, and the two strings differ.
1224        validate_marginal_slope_z_column_exclusion(
1225            &main,
1226            &parse_formula("y ~ 1").expect("parse logslope"),
1227            "pgs_ctn_z",
1228            "bernoulli marginal-slope",
1229            "logslope_formula",
1230        )
1231        .expect("literal-name validator cannot see the alias");
1232
1233        let err = validate_marginal_slope_z_alias_exclusion(
1234            &main,
1235            &cols,
1236            "pgs_ctn_z",
1237            "bernoulli marginal-slope",
1238        )
1239        .expect_err("bare `z` resolves to the reserved score column");
1240        assert!(err.contains("reserves z column 'pgs_ctn_z'"), "{err}");
1241        assert!(err.contains("cannot also appear in the main formula"), "{err}");
1242    }
1243
1244    /// Negative controls: the alias guard must not reject anything else.
1245    #[test]
1246    fn marginal_slope_alias_validator_accepts_legitimate_formulas() {
1247        // A baseline that never mentions the score, under either spelling.
1248        validate_marginal_slope_z_alias_exclusion(
1249            &parse_formula("y ~ x + w").expect("parse clean main"),
1250            &frame(&["y", "x", "w", "pgs_ctn_z"]),
1251            "pgs_ctn_z",
1252            "bernoulli marginal-slope",
1253        )
1254        .expect("a baseline without the score is legitimate");
1255
1256        // A frame carrying its OWN `z` column keeps it: `column_map_with_alias`
1257        // inserts with `or_insert`, so the alias is inert and `z` in the main
1258        // formula means that column, not the score.
1259        validate_marginal_slope_z_alias_exclusion(
1260            &parse_formula("y ~ x + z").expect("parse main with real z"),
1261            &frame(&["y", "x", "z", "pgs_ctn_z"]),
1262            "pgs_ctn_z",
1263            "bernoulli marginal-slope",
1264        )
1265        .expect("a real `z` column is not the alias");
1266
1267        // z_column == "z": the literal-name validator already covers this
1268        // verbatim, so the alias guard must stay out of its way.
1269        validate_marginal_slope_z_alias_exclusion(
1270            &parse_formula("y ~ x + z").expect("parse main"),
1271            &frame(&["y", "x", "z"]),
1272            "z",
1273            "bernoulli marginal-slope",
1274        )
1275        .expect("alias guard defers to the literal-name validator when they coincide");
1276    }
1277
1278    /// The alias guard is an addition, not a replacement: spelling the reserved
1279    /// column out in full is still rejected, by the original validator.
1280    #[test]
1281    fn marginal_slope_literal_z_column_is_still_rejected() {
1282        let err = validate_marginal_slope_z_column_exclusion(
1283            &parse_formula("y ~ x + pgs_ctn_z").expect("parse main"),
1284            &parse_formula("y ~ 1").expect("parse logslope"),
1285            "pgs_ctn_z",
1286            "bernoulli marginal-slope",
1287            "logslope_formula",
1288        )
1289        .expect_err("the reserved column named in full must still be rejected");
1290        assert!(err.contains("reserves z column 'pgs_ctn_z'"), "{err}");
1291    }
1292
1293    #[test]
1294    fn logslope_surface_declarations_are_additive() {
1295        let parsed = parse_formula("y ~ s(pc1) + logslope(z2, s(pc2)) + logslope(z3, x3)")
1296            .expect("parse additive logslope surfaces");
1297        assert_eq!(parsed.terms.len(), 1);
1298        assert_eq!(parsed.logslope_surfaces.len(), 2);
1299        assert_eq!(parsed.logslope_surfaces[0].z_column, "z2");
1300        assert_eq!(parsed.logslope_surfaces[0].terms.len(), 1);
1301        assert_eq!(parsed.logslope_surfaces[1].z_column, "z3");
1302        assert_eq!(parsed.logslope_surfaces[1].terms.len(), 1);
1303    }
1304
1305    #[test]
1306    fn marginal_slope_z_column_validator_reserves_all_surface_z_columns() {
1307        let main = parse_formula("y ~ x").expect("parse main");
1308        let logslope = parse_formula("y ~ s(pc1) + logslope(z2, s(z3)) + logslope(z3, x)")
1309            .expect("parse logslope surfaces");
1310        let err = validate_marginal_slope_z_column_exclusion(
1311            &main,
1312            &logslope,
1313            "z1",
1314            "bernoulli marginal-slope",
1315            "--logslope-formula",
1316        )
1317        .expect_err("surface formula should reject another reserved z coordinate");
1318        assert!(err.contains("reserves z column 'z3'"));
1319    }
1320
1321    /// Extract the single `RandomEffect` term's unseen-level policy from a
1322    /// one-term formula, panicking if the term is not a random-effect block.
1323    fn random_effect_lenient_unseen(formula: &str) -> bool {
1324        let parsed = parse_formula(formula).expect("parse random-effect formula");
1325        let re = parsed.terms.iter().find_map(|t| match t {
1326            ParsedTerm::RandomEffect { lenient_unseen, .. } => Some(*lenient_unseen),
1327            _ => None,
1328        });
1329        re.unwrap_or_else(|| panic!("{formula} did not lower to a RandomEffect term"))
1330    }
1331
1332    #[test]
1333    fn factor_wrapper_is_strict_on_unseen_levels_while_group_re_are_lenient() {
1334        // Regression for #2137 (sibling of #2102): `factor(g)` is a FIXED
1335        // categorical factor (R `factor()` / patsy `C()`), so an out-of-vocabulary
1336        // level at predict is a schema mismatch that must raise — NOT be shrunk to
1337        // the centering point. `group(g)`/`re(g)`/`s(g, bs="re")` are genuine
1338        // random effects that tolerate a held-out group (→ population mean). The
1339        // parse arm once hardcoded `lenient_unseen: true` for all four wrappers,
1340        // so `factor(g)` silently averaged an unseen level. Pin the per-wrapper
1341        // policy at the parse layer, where the whole distinction now lives.
1342        assert!(
1343            !random_effect_lenient_unseen("y ~ factor(g)"),
1344            "factor(g) is a fixed categorical factor: strict (lenient_unseen=false) on unseen levels"
1345        );
1346        for lenient in ["y ~ group(g)", "y ~ re(g)", "y ~ s(g, bs=re)"] {
1347            assert!(
1348                random_effect_lenient_unseen(lenient),
1349                "{lenient} is a genuine random effect: lenient (lenient_unseen=true) on unseen levels"
1350            );
1351        }
1352    }
1353}
1354
1355// ---------------------------------------------------------------------------
1356// Higher-level formula parsing: ParsedFormula, ParsedTerm, and friends
1357// ---------------------------------------------------------------------------
1358
1359#[derive(Clone, Debug)]
1360pub struct LinkWiggleFormulaSpec {
1361    pub degree: usize,
1362    pub num_internal_knots: usize,
1363    pub penalty_orders: Vec<usize>,
1364    pub double_penalty: bool,
1365}
1366
1367pub fn default_linkwiggle_formulaspec() -> LinkWiggleFormulaSpec {
1368    let cfg = WigglePenaltyConfig::cubic_triple_operator_default();
1369    LinkWiggleFormulaSpec {
1370        degree: cfg.degree,
1371        num_internal_knots: cfg.num_internal_knots,
1372        penalty_orders: cfg.penalty_orders,
1373        double_penalty: cfg.double_penalty,
1374    }
1375}
1376
1377#[derive(Clone, Debug)]
1378pub struct LinkFormulaSpec {
1379    pub link: String,
1380    pub mixture_rho: Option<String>,
1381    pub sas_init: Option<String>,
1382    pub beta_logistic_init: Option<String>,
1383}
1384
1385#[derive(Clone, Debug)]
1386pub struct SurvivalFormulaSpec {
1387    pub spec: Option<String>,
1388    pub survival_distribution: Option<String>,
1389}
1390
1391#[derive(Clone, Debug)]
1392pub struct ParsedFormula {
1393    pub response: String,
1394    pub terms: Vec<ParsedTerm>,
1395    pub logslope_surfaces: Vec<LogSlopeSurfaceSpec>,
1396    pub linkwiggle: Option<LinkWiggleFormulaSpec>,
1397    pub timewiggle: Option<LinkWiggleFormulaSpec>,
1398    pub linkspec: Option<LinkFormulaSpec>,
1399    pub survivalspec: Option<SurvivalFormulaSpec>,
1400}
1401
1402#[derive(Clone, Debug)]
1403pub struct LogSlopeSurfaceSpec {
1404    pub z_column: String,
1405    pub terms: Vec<ParsedTerm>,
1406}
1407
1408pub fn marginal_slope_logslope_surfaces(
1409    logslope_formula: &ParsedFormula,
1410    default_z_column: &str,
1411) -> Result<Vec<LogSlopeSurfaceSpec>, String> {
1412    let mut surfaces = Vec::new();
1413    if !logslope_formula.terms.is_empty() {
1414        surfaces.push(LogSlopeSurfaceSpec {
1415            z_column: default_z_column.to_string(),
1416            terms: logslope_formula.terms.clone(),
1417        });
1418    }
1419    surfaces.extend(logslope_formula.logslope_surfaces.clone());
1420    if surfaces.is_empty() {
1421        surfaces.push(LogSlopeSurfaceSpec {
1422            z_column: default_z_column.to_string(),
1423            terms: Vec::new(),
1424        });
1425    }
1426    let mut seen = std::collections::BTreeSet::<String>::new();
1427    for surface in &surfaces {
1428        if !seen.insert(surface.z_column.clone()) {
1429            return Err(FormulaDslError::IncompatibleTerm {
1430                reason: format!(
1431                    "logslope formula declares z column '{}' more than once; each z coordinate needs exactly one log-slope surface",
1432                    surface.z_column
1433                ),
1434            }
1435            .into());
1436        }
1437    }
1438    Ok(surfaces)
1439}
1440
1441#[derive(Clone, Debug)]
1442pub enum ParsedTerm {
1443    Linear {
1444        name: String,
1445        explicit: bool,
1446        double_penalty: bool,
1447        coefficient_min: Option<f64>,
1448        coefficient_max: Option<f64>,
1449    },
1450    BoundedLinear {
1451        name: String,
1452        min: f64,
1453        max: f64,
1454        prior: BoundedCoefficientPriorSpec,
1455        double_penalty: bool,
1456    },
1457    RandomEffect {
1458        name: String,
1459        /// Unseen-level policy, fixed at parse time by the wrapper the user
1460        /// wrote. `group(g)`/`re(g)`/`s(g, bs="re")` are genuine **random
1461        /// effects**: a held-out group is shrunk to the population mean, so an
1462        /// unseen level at predict is tolerated (`true`). `factor(g)` is a
1463        /// **fixed** categorical factor (R `factor()` / patsy `C()`
1464        /// convention): like a bare `+ g` categorical main effect, an unseen
1465        /// level is a schema mismatch that must raise rather than collapse onto
1466        /// the factor's centering point (`false`, #2137/#2102). Both wrappers
1467        /// share the penalized-categorical materialization; only this policy
1468        /// distinguishes them, so seen-level fits are identical.
1469        lenient_unseen: bool,
1470    },
1471    Smooth {
1472        label: String,
1473        vars: Vec<String>,
1474        kind: SmoothKind,
1475        options: BTreeMap<String, String>,
1476    },
1477    LinkWiggle {
1478        options: BTreeMap<String, String>,
1479    },
1480    TimeWiggle {
1481        options: BTreeMap<String, String>,
1482    },
1483    LinkConfig {
1484        options: BTreeMap<String, String>,
1485    },
1486    SurvivalConfig {
1487        options: BTreeMap<String, String>,
1488    },
1489    LogSlopeSurface {
1490        z_column: String,
1491        terms: Vec<ParsedTerm>,
1492    },
1493    /// Wilkinson-Rogers interaction term `a:b[:c...]`.
1494    ///
1495    /// `vars` is a sorted, deduplicated list of base column names. Each element
1496    /// must be a bare identifier — interactions with function-call atoms
1497    /// (smooths, factors, etc.) are rejected upstream because their design
1498    /// columns are not simple products. The design column is the elementwise
1499    /// product of the referenced numeric columns.
1500    Interaction {
1501        vars: Vec<String>,
1502        double_penalty: bool,
1503    },
1504}
1505
1506/// Collect the names of every data column the parsed terms consume.
1507///
1508/// This is the canonical formula→columns walk shared by the fit-time and
1509/// predict-time required-column computations (the CLI and PyFFI surfaces both
1510/// route through it). It includes a smooth's positional `vars` *and* its `by=`
1511/// grouping/scaling column (`s(x, by=g)`), which `term_builder` reads from
1512/// `options["by"]` but which is not among the positional variables — omitting
1513/// it would drop a genuine predictor from the model's input contract.
1514pub fn parsed_term_column_names(
1515    terms: &[ParsedTerm],
1516    out: &mut std::collections::BTreeSet<String>,
1517) {
1518    for term in terms {
1519        match term {
1520            ParsedTerm::Linear { name, .. }
1521            | ParsedTerm::BoundedLinear { name, .. }
1522            | ParsedTerm::RandomEffect { name, .. } => {
1523                out.insert(name.clone());
1524            }
1525            ParsedTerm::Smooth { vars, options, .. } => {
1526                out.extend(vars.iter().cloned());
1527                if let Some(by) = options.get("by") {
1528                    out.insert(by.clone());
1529                }
1530            }
1531            ParsedTerm::Interaction { vars, .. } => {
1532                out.extend(vars.iter().cloned());
1533            }
1534            ParsedTerm::LinkWiggle { .. }
1535            | ParsedTerm::TimeWiggle { .. }
1536            | ParsedTerm::LinkConfig { .. }
1537            | ParsedTerm::SurvivalConfig { .. } => {}
1538            ParsedTerm::LogSlopeSurface { z_column, terms } => {
1539                out.insert(z_column.clone());
1540                parsed_term_column_names(terms, out);
1541            }
1542        }
1543    }
1544}
1545
1546pub fn parsed_terms_reference_column(terms: &[ParsedTerm], column_name: &str) -> bool {
1547    terms.iter().any(|term| match term {
1548        ParsedTerm::Linear { name, .. }
1549        | ParsedTerm::BoundedLinear { name, .. }
1550        | ParsedTerm::RandomEffect { name, .. } => name == column_name,
1551        ParsedTerm::Smooth { vars, options, .. } => {
1552            vars.iter().any(|var| var == column_name)
1553                || options.get("by").is_some_and(|by| by == column_name)
1554        }
1555        ParsedTerm::Interaction { vars, .. } => vars.iter().any(|var| var == column_name),
1556        ParsedTerm::LinkWiggle { .. }
1557        | ParsedTerm::TimeWiggle { .. }
1558        | ParsedTerm::LinkConfig { .. }
1559        | ParsedTerm::SurvivalConfig { .. } => false,
1560        ParsedTerm::LogSlopeSurface { z_column, terms } => {
1561            z_column == column_name || parsed_terms_reference_column(terms, column_name)
1562        }
1563    })
1564}
1565
1566/// Reject the score entering the main formula through its canonical alias.
1567///
1568/// `validate_marginal_slope_z_column_exclusion` compares each parsed term's
1569/// literal name against the configured `z_column`, so it rejects
1570/// `disease ~ pgs_ctn_z` — and accepts `disease ~ z`, which
1571/// `column_map_with_alias(col_map, "z", z_column)` then resolves to *the same
1572/// column*. The score lands in the marginal design either way; only one
1573/// spelling is caught. That is not a cosmetic gap: with an intercept-only
1574/// log-slope the effective log-slope direction is `f = a*1 + s*z`, which lies
1575/// in the span of the effective marginal design `[1, z]` row by row for ANY
1576/// data, so the fit is degenerate by construction. It used to surface ~12 s
1577/// later as a `FullyConfounded` refusal from the BMS audit, a message that
1578/// reads like a solver failure rather than a formula mistake (gam#2432).
1579///
1580/// Checked against the frame rather than the formula text because the alias is
1581/// only live when the frame has no real `z` column of its own — see
1582/// [`marginal_slope_z_alias_is_live`]. A frame that genuinely carries `z` keeps
1583/// it, and writing `z` there is legitimate and still accepted.
1584pub fn validate_marginal_slope_z_alias_exclusion(
1585    main_formula: &ParsedFormula,
1586    col_map: &HashMap<String, usize>,
1587    z_column: &str,
1588    context: &str,
1589) -> Result<(), String> {
1590    if z_column == MARGINAL_SLOPE_Z_ALIAS {
1591        // Already covered verbatim by the literal-name exclusion.
1592        return Ok(());
1593    }
1594    if !marginal_slope_z_alias_is_live(col_map, z_column) {
1595        return Ok(());
1596    }
1597    if parsed_terms_reference_column(&main_formula.terms, MARGINAL_SLOPE_Z_ALIAS) {
1598        return Err(FormulaDslError::IncompatibleTerm {
1599            reason: format!(
1600                "{context} reserves z column '{z_column}' as the auxiliary latent score, and the \
1601                 bare name '{MARGINAL_SLOPE_Z_ALIAS}' in the main formula resolves to it; it \
1602                 cannot also appear in the main formula. Give the baseline its own covariates \
1603                 (the score's effect is carried by the log-slope formula)."
1604            ),
1605        }
1606        .into());
1607    }
1608    Ok(())
1609}
1610
1611pub fn validate_marginal_slope_z_column_exclusion(
1612    main_formula: &ParsedFormula,
1613    logslope_formula: &ParsedFormula,
1614    z_column: &str,
1615    context: &str,
1616    logslope_label: &str,
1617) -> Result<(), String> {
1618    let surfaces = marginal_slope_logslope_surfaces(logslope_formula, z_column)?;
1619    // The CLI/configured z column is reserved even when the log-slope formula
1620    // is intercept-only (`~ 1`) and therefore contributes no surface terms.
1621    // Explicit logslope(...) declarations may reserve additional z coordinates.
1622    let mut reserved_z_columns = std::collections::BTreeSet::<&str>::new();
1623    reserved_z_columns.insert(z_column);
1624    reserved_z_columns.extend(surfaces.iter().map(|surface| surface.z_column.as_str()));
1625
1626    for reserved in &reserved_z_columns {
1627        if parsed_terms_reference_column(&main_formula.terms, reserved) {
1628            return Err(FormulaDslError::IncompatibleTerm {
1629                reason: format!(
1630                    "{context} reserves z column '{reserved}' as the auxiliary latent score; it cannot also appear in the main formula"
1631                ),
1632            }
1633            .into());
1634        }
1635    }
1636    for reserved in &reserved_z_columns {
1637        if parsed_terms_reference_column(&logslope_formula.terms, reserved) {
1638            return Err(FormulaDslError::IncompatibleTerm {
1639                reason: format!(
1640                    "{context} reserves z column '{reserved}' as the auxiliary latent score; it cannot also appear in {logslope_label}"
1641                ),
1642            }
1643            .into());
1644        }
1645        for surface in &surfaces {
1646            if parsed_terms_reference_column(&surface.terms, reserved) {
1647                return Err(FormulaDslError::IncompatibleTerm {
1648                    reason: format!(
1649                        "{context} reserves z column '{reserved}' as an auxiliary latent score; it cannot also appear in {logslope_label}"
1650                    ),
1651                }
1652                .into());
1653            }
1654        }
1655    }
1656    Ok(())
1657}
1658
1659#[derive(Clone, Copy, Debug)]
1660pub enum SmoothKind {
1661    S,
1662    Te,
1663    /// Tensor smooth (`t2(...)`) using mgcv's separable penalty decomposition:
1664    /// the tensor coefficient space is split into marginal penalized/null-space
1665    /// tensor subspaces, with one smoothing parameter per non-null subspace.
1666    T2,
1667    /// Tensor *interaction* smooth (`ti(...)`): a tensor-product smooth whose
1668    /// marginal main effects are excluded, so the term captures only the pure
1669    /// interaction between its variables. Materializes through the same tensor
1670    /// path as [`SmoothKind::Te`] but with per-margin sum-to-zero
1671    /// identifiability (`TensorBSplineIdentifiability::MarginalSumToZero`).
1672    Ti,
1673}
1674
1675#[derive(Clone, Copy, Debug)]
1676pub enum LinkMode {
1677    Strict,
1678    Flexible,
1679}
1680
1681#[derive(Clone, Debug)]
1682pub struct LinkChoice {
1683    pub mode: LinkMode,
1684    pub link: LinkFunction,
1685    pub mixture_components: Option<Vec<LinkComponent>>,
1686}
1687
1688// ---------------------------------------------------------------------------
1689// Link wiggle / link choice helpers
1690// ---------------------------------------------------------------------------
1691
1692pub fn effectivelinkwiggle_formulaspec(
1693    formula_linkwiggle: Option<&LinkWiggleFormulaSpec>,
1694    link_choice: Option<&LinkChoice>,
1695) -> Option<LinkWiggleFormulaSpec> {
1696    formula_linkwiggle.cloned().or_else(|| {
1697        link_choice.and_then(|choice| {
1698            if matches!(choice.mode, LinkMode::Flexible) {
1699                Some(default_linkwiggle_formulaspec())
1700            } else {
1701                None
1702            }
1703        })
1704    })
1705}
1706
1707pub const fn linkname_supports_joint_wiggle(link: LinkFunction) -> bool {
1708    !matches!(link, LinkFunction::Sas | LinkFunction::BetaLogistic)
1709}
1710
1711pub const fn linkchoice_supports_joint_wiggle(choice: &LinkChoice) -> bool {
1712    match &choice.mixture_components {
1713        None => linkname_supports_joint_wiggle(choice.link),
1714        Some(_) => false,
1715    }
1716}
1717
1718pub fn require_linkchoice_supports_joint_wiggle(
1719    choice: &LinkChoice,
1720    context: &str,
1721) -> Result<(), String> {
1722    if linkchoice_supports_joint_wiggle(choice) {
1723        Ok(())
1724    } else {
1725        Err(joint_wiggle_unsupported_link_message(context))
1726    }
1727}
1728
1729pub const fn likelihood_spec_supports_joint_wiggle(likelihood: &LikelihoodSpec) -> bool {
1730    inverse_link_supports_joint_wiggle(&likelihood.link)
1731}
1732
1733pub fn require_likelihood_spec_supports_joint_wiggle(
1734    likelihood: &LikelihoodSpec,
1735    context: &str,
1736) -> Result<(), String> {
1737    if likelihood_spec_supports_joint_wiggle(likelihood) {
1738        Ok(())
1739    } else {
1740        Err(joint_wiggle_unsupported_link_message(context))
1741    }
1742}
1743
1744/// Family-agnostic capability of the joint link-wiggle machinery: which base
1745/// inverse links a monotone warp can be fit over AND reconstructed from at
1746/// predict time. Every state-less standard link qualifies — the warp fit and
1747/// its saved-model reconstruction evaluate the base inverse link purely through
1748/// the generic `inverse_link_jet_for_inverse_link` jet dispatch, which carries
1749/// LogLog and Cauchit exactly as it does Logit/Probit/CLogLog. LogLog/Cauchit
1750/// were previously omitted, so a binomial `flexible(loglog)`/`flexible(cauchit)`
1751/// fit that this-gate-agnostically *converged* (see
1752/// `binomial_inverse_link_supports_joint_wiggle`) then failed at predict when
1753/// `FittedModel::saved_link_wiggle` re-checked the saved link here (#2155). The
1754/// state-bearing links (SAS/BetaLogistic/Mixture/LatentCLogLog) carry fitted
1755/// warp/skew state of their own and are intentionally excluded.
1756pub const fn inverse_link_supports_joint_wiggle(link: &InverseLink) -> bool {
1757    matches!(
1758        link,
1759        InverseLink::Standard(StandardLink::Identity)
1760            | InverseLink::Standard(StandardLink::Log)
1761            | InverseLink::Standard(StandardLink::Logit)
1762            | InverseLink::Standard(StandardLink::Probit)
1763            | InverseLink::Standard(StandardLink::CLogLog)
1764            | InverseLink::Standard(StandardLink::LogLog)
1765            | InverseLink::Standard(StandardLink::Cauchit)
1766    )
1767}
1768
1769pub fn require_inverse_link_supports_joint_wiggle(
1770    link: &InverseLink,
1771    context: &str,
1772) -> Result<(), String> {
1773    if inverse_link_supports_joint_wiggle(link) {
1774        Ok(())
1775    } else {
1776        Err(joint_wiggle_unsupported_link_message(context))
1777    }
1778}
1779
1780/// Which binomial base links the joint link-wiggle (flexible-link) solver can
1781/// fit. All five standard binomial probability links qualify: the wiggle kernel
1782/// (`BinomialMeanWiggleFamily`) evaluates the base inverse link purely through
1783/// the generic `inverse_link_jet_for_inverse_link` dispatch, which carries full
1784/// jets for LogLog and Cauchit exactly as it does for Logit/Probit/CLogLog — so
1785/// `flexible(loglog)` / `flexible(cauchit)` fit through the same machinery
1786/// (#2155). Previously this gate listed only logit/probit/cloglog while the
1787/// permissive parse gate `linkname_supports_joint_wiggle` admitted loglog/cauchit,
1788/// so the config was accepted then aborted deep in the solver. The state-bearing
1789/// links (SAS/BetaLogistic/Mixture/LatentCLogLog) and identity/log stay out: the
1790/// warp is defined only over a fixed state-less base probability link.
1791pub const fn binomial_inverse_link_supports_joint_wiggle(link: &InverseLink) -> bool {
1792    matches!(
1793        link,
1794        InverseLink::Standard(StandardLink::Logit)
1795            | InverseLink::Standard(StandardLink::Probit)
1796            | InverseLink::Standard(StandardLink::CLogLog)
1797            | InverseLink::Standard(StandardLink::LogLog)
1798            | InverseLink::Standard(StandardLink::Cauchit)
1799    )
1800}
1801
1802pub fn require_binomial_inverse_link_supports_joint_wiggle(
1803    link: &InverseLink,
1804    context: &str,
1805) -> Result<(), String> {
1806    if binomial_inverse_link_supports_joint_wiggle(link) {
1807        Ok(())
1808    } else {
1809        Err(FormulaDslError::IncompatibleTerm {
1810            reason: format!(
1811                "{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)"
1812            ),
1813        }
1814        .into())
1815    }
1816}
1817
1818pub fn joint_wiggle_unsupported_link_message(context: &str) -> String {
1819    format!(
1820        "{context} does not support latent-cloglog, SAS, BetaLogistic, or Mixture links; wiggle is only available for jointly fitted standard links"
1821    )
1822}
1823
1824// ---------------------------------------------------------------------------
1825// Option-map helpers (shared by formula parsing and term construction)
1826// ---------------------------------------------------------------------------
1827
1828pub fn option_usize(map: &BTreeMap<String, String>, key: &str) -> Option<usize> {
1829    map.get(key).and_then(|v| v.parse::<usize>().ok())
1830}
1831
1832/// Local sibling of `term_builder::validate_known_options` used by the
1833/// parser-side `linear / bounded / constrain / nonnegative / nonpositive`
1834/// branches (which build their `ParsedTerm` here and never enter
1835/// `term_builder::build_smooth_basis`). Without this, typos like
1836/// `bounded(x, min=0, max=1, foo=bar)` silently succeed because the
1837/// `foo` key was just never read.
1838fn validate_known_term_options(
1839    term_name: &str,
1840    options: &BTreeMap<String, String>,
1841    known: &[&str],
1842    raw: &str,
1843) -> Result<(), String> {
1844    let known_set: std::collections::BTreeSet<&&str> = known.iter().collect();
1845    for key in options.keys() {
1846        if !known_set.contains(&key.as_str()) {
1847            let known_sorted = {
1848                let mut v = known.to_vec();
1849                v.sort_unstable();
1850                v.join(", ")
1851            };
1852            let known_hint = if known.is_empty() {
1853                "no options".to_string()
1854            } else {
1855                format!("[{known_sorted}]")
1856            };
1857            return Err(FormulaDslError::InvalidArgument {
1858                reason: format!(
1859                    "{term_name}() does not accept option `{key}` (in `{raw}`); known options: {known_hint}"
1860                ),
1861            }
1862            .into());
1863        }
1864    }
1865    Ok(())
1866}
1867
1868pub fn option_usize_any(map: &BTreeMap<String, String>, keys: &[&str]) -> Option<usize> {
1869    for key in keys {
1870        if let Some(v) = option_usize(map, key) {
1871            return Some(v);
1872        }
1873    }
1874    None
1875}
1876
1877/// Strict integer option: returns `Ok(None)` if not present, `Ok(Some(n))` if
1878/// it parses as a non-negative integer, and `Err(msg)` if the user supplied a
1879/// value that isn't a valid usize (negative, decimal, garbage). Without this
1880/// the lenient `option_usize` silently drops invalid values and reverts to
1881/// the default — `k=-1` and `k=1.5` were both accepted as "k not specified"
1882/// instead of being flagged as user mistakes.
1883pub fn option_usize_strict(
1884    map: &BTreeMap<String, String>,
1885    key: &str,
1886) -> Result<Option<usize>, String> {
1887    match map.get(key) {
1888        None => Ok(None),
1889        Some(raw) => raw.parse::<usize>().map(Some).map_err(|err| {
1890            FormulaDslError::InvalidArgument {
1891                reason: format!(
1892                    "option `{key}={raw}` is not a non-negative integer; \
1893                     expected a whole number >= 0: {err}"
1894                ),
1895            }
1896            .into()
1897        }),
1898    }
1899}
1900
1901/// Strict variant of `option_usize_any` that errors on the first present-but-
1902/// unparseable key rather than silently falling through.
1903pub fn option_usize_any_strict(
1904    map: &BTreeMap<String, String>,
1905    keys: &[&str],
1906) -> Result<Option<usize>, String> {
1907    for key in keys {
1908        if let Some(v) = option_usize_strict(map, key)? {
1909            return Ok(Some(v));
1910        }
1911    }
1912    Ok(None)
1913}
1914
1915pub fn option_f64(map: &BTreeMap<String, String>, key: &str) -> Option<f64> {
1916    map.get(key).and_then(|v| v.parse::<f64>().ok())
1917}
1918
1919/// Strict float option: `Ok(None)` if absent, `Ok(Some(n))` if parses as a
1920/// finite f64, `Err` if the user passed an unparseable value (rather than
1921/// silently dropping it like the lenient `option_f64`).
1922pub fn option_f64_strict(map: &BTreeMap<String, String>, key: &str) -> Result<Option<f64>, String> {
1923    match map.get(key) {
1924        None => Ok(None),
1925        Some(raw) => match raw.parse::<f64>() {
1926            Ok(v) if v.is_finite() => Ok(Some(v)),
1927            Ok(v) => Err(FormulaDslError::InvalidArgument {
1928                reason: format!("option `{key}={raw}` parses as {v} which is not a finite number"),
1929            }
1930            .into()),
1931            Err(err) => Err(FormulaDslError::InvalidArgument {
1932                reason: format!(
1933                    "option `{key}={raw}` is not a valid number; expected a finite decimal: {err}"
1934                ),
1935            }
1936            .into()),
1937        },
1938    }
1939}
1940
1941pub fn option_bool(map: &BTreeMap<String, String>, key: &str) -> Option<bool> {
1942    map.get(key)
1943        .and_then(|v| match v.trim().to_ascii_lowercase().as_str() {
1944            "true" | "1" | "yes" | "y" => Some(true),
1945            "false" | "0" | "no" | "n" => Some(false),
1946            _ => None,
1947        })
1948}
1949
1950/// Strict boolean option: `Ok(None)` if absent, `Ok(Some(b))` for a recognized
1951/// truthy/falsy token, and `Err(msg)` for a present-but-unparseable value. The
1952/// lenient `option_bool` maps an unrecognized value to `None`, which callers
1953/// then silently treat as "not specified" — masking user typos like
1954/// `double_penalty=ture`.
1955pub fn option_bool_strict(
1956    map: &BTreeMap<String, String>,
1957    key: &str,
1958) -> Result<Option<bool>, String> {
1959    match map.get(key) {
1960        None => Ok(None),
1961        Some(raw) => match raw.trim().to_ascii_lowercase().as_str() {
1962            "true" | "1" | "yes" | "y" => Ok(Some(true)),
1963            "false" | "0" | "no" | "n" => Ok(Some(false)),
1964            _ => Err(FormulaDslError::InvalidArgument {
1965                reason: format!(
1966                    "option `{key}={raw}` is not a boolean; \
1967                     expected one of true/false/yes/no/1/0"
1968                ),
1969            }
1970            .into()),
1971        },
1972    }
1973}
1974
1975pub fn strip_quotes(v: &str) -> &str {
1976    let b = v.as_bytes();
1977    if b.len() >= 2
1978        && ((b[0] == b'\'' && b[b.len() - 1] == b'\'') || (b[0] == b'"' && b[b.len() - 1] == b'"'))
1979    {
1980        &v[1..v.len() - 1]
1981    } else {
1982        v
1983    }
1984}
1985
1986// ---------------------------------------------------------------------------
1987// Sub-parsers for formula option blocks
1988// ---------------------------------------------------------------------------
1989
1990fn parse_linear_constraint_bounds(
1991    options: &BTreeMap<String, String>,
1992    raw: &str,
1993) -> Result<(Option<f64>, Option<f64>), String> {
1994    let min = parse_optional_f64_option_alias(options, &["min", "lower"], raw, "linear")?;
1995    let max = parse_optional_f64_option_alias(options, &["max", "upper"], raw, "linear")?;
1996    if let (Some(min), Some(max)) = (min, max)
1997        && (!min.is_finite() || !max.is_finite() || min > max)
1998    {
1999        return Err(FormulaDslError::InvalidArgument {
2000            reason: format!(
2001                "linear coefficient constraints require finite min <= max, got min={min}, max={max}: {raw}"
2002            ),
2003        }
2004        .into());
2005    }
2006    Ok((min, max))
2007}
2008
2009fn parse_required_f64_option(
2010    options: &BTreeMap<String, String>,
2011    key: &str,
2012    raw: &str,
2013) -> Result<f64, String> {
2014    let value = options
2015        .get(key)
2016        .ok_or_else(|| FormulaDslError::MalformedConfig {
2017            reason: format!("bounded() is missing required '{key}' argument: {raw}"),
2018        })?;
2019    value.parse::<f64>().map_err(|err| {
2020        FormulaDslError::InvalidArgument {
2021            reason: format!(
2022                "bounded() argument '{key}' must be a finite number, got '{}': {err}: {raw}",
2023                value
2024            ),
2025        }
2026        .into()
2027    })
2028}
2029
2030fn parse_optional_f64_option(
2031    options: &BTreeMap<String, String>,
2032    key: &str,
2033    raw: &str,
2034) -> Result<Option<f64>, String> {
2035    match options.get(key) {
2036        Some(value) => value.parse::<f64>().map(Some).map_err(|err| {
2037            FormulaDslError::InvalidArgument {
2038                reason: format!(
2039                    "bounded() argument '{key}' must be a finite number, got '{}': {err}: {raw}",
2040                    value
2041                ),
2042            }
2043            .into()
2044        }),
2045        None => Ok(None),
2046    }
2047}
2048
2049fn parse_optional_f64_option_alias(
2050    options: &BTreeMap<String, String>,
2051    keys: &[&str],
2052    raw: &str,
2053    fn_label: &str,
2054) -> Result<Option<f64>, String> {
2055    let mut found: Option<(&str, f64)> = None;
2056    for key in keys {
2057        if let Some(value) = options.get(*key) {
2058            let parsed = value
2059                .parse::<f64>()
2060                .map_err(|err| FormulaDslError::InvalidArgument {
2061                    reason: format!(
2062                        "{fn_label}() argument '{key}' must be a finite number, got '{}': {err}: {raw}",
2063                        value
2064                    ),
2065                })?;
2066            if found.is_some() {
2067                return Err(FormulaDslError::IncompatibleTerm {
2068                    reason: format!(
2069                        "{fn_label}() cannot specify both '{}' and '{}': {raw}",
2070                        found.expect("present").0,
2071                        key
2072                    ),
2073                }
2074                .into());
2075            }
2076            found = Some((key, parsed));
2077        }
2078    }
2079    Ok(found.map(|(_, v)| v))
2080}
2081
2082fn parse_linkwiggle_penalty_orders(raw: Option<&str>) -> Result<Vec<usize>, String> {
2083    let Some(raw) = raw.map(str::trim) else {
2084        return Ok(WigglePenaltyConfig::cubic_triple_operator_default().penalty_orders);
2085    };
2086    if raw.is_empty() {
2087        return Ok(WigglePenaltyConfig::cubic_triple_operator_default().penalty_orders);
2088    }
2089    let mut out = Vec::<usize>::new();
2090    for token in raw.split(',') {
2091        let t = token.trim().to_ascii_lowercase();
2092        if t.is_empty() {
2093            continue;
2094        }
2095        match t.as_str() {
2096            "all" => {
2097                out.extend([1, 2, 3]);
2098            }
2099            "slope" | "1" => out.push(1),
2100            "curvature" | "2" => out.push(2),
2101            "curvature-change" | "curvature_change" | "3" => out.push(3),
2102            _ => {
2103                return Err(FormulaDslError::InvalidArgument {
2104                    reason: format!(
2105                        "invalid linkwiggle penalty_order '{t}'; use all|slope|curvature|curvature-change or 1/2/3"
2106                    ),
2107                }
2108                .into());
2109            }
2110        }
2111    }
2112    if out.is_empty() {
2113        out.extend(WigglePenaltyConfig::cubic_triple_operator_default().penalty_orders);
2114    }
2115    out.sort_unstable();
2116    out.dedup();
2117    Ok(out)
2118}
2119
2120pub fn parse_linkwiggle_formulaspec(
2121    options: &BTreeMap<String, String>,
2122    raw: &str,
2123) -> Result<LinkWiggleFormulaSpec, String> {
2124    let allowed = [
2125        "degree",
2126        "internal_knots",
2127        "penalty_order",
2128        "double_penalty",
2129    ];
2130    let unknown = options
2131        .keys()
2132        .filter(|key| !allowed.contains(&key.as_str()))
2133        .cloned()
2134        .collect::<Vec<_>>();
2135    let term_name = raw.split('(').next().unwrap_or("linkwiggle");
2136    if !unknown.is_empty() {
2137        return Err(FormulaDslError::InvalidArgument {
2138            reason: format!(
2139                "{}() does not support option(s) {}: {raw}",
2140                term_name,
2141                unknown.join(", ")
2142            ),
2143        }
2144        .into());
2145    }
2146    let defaults = WigglePenaltyConfig::cubic_triple_operator_default();
2147    // Strict parsing: a present-but-unparseable value (`degree=abc`, `=-3`,
2148    // `=6.5`) must be rejected, not silently dropped and replaced by the
2149    // default as the lossy `option_usize`/`option_bool` readers would do.
2150    //
2151    // This parser is shared by *all* wiggle grammars: `linkwiggle` and
2152    // `timewiggle` (see `parse_formula`), the standard-model flexible-link
2153    // wiggle, and the marginal-slope score-warp / link-deviation routing.
2154    // The general monotone I-spline value basis (`monotone_wiggle_*` in
2155    // `families::gamlss`, used by `timewiggle` and the location-scale survival
2156    // path) honors arbitrary `degree >= 2`, while only the cubic-only
2157    // score-warp / link-deviation `DeviationRuntime` is restricted to 3.
2158    // A consumer-specific limit therefore must NOT be baked into this shared
2159    // parser — it is enforced at the routing layer that feeds the cubic-only
2160    // runtime (`deviation_block_config_from_formula_linkwiggle`). Here we only
2161    // enforce the universal lower bound that a polynomial degree is positive.
2162    let degree = option_usize_strict(options, "degree")?.unwrap_or(defaults.degree);
2163    if degree < 1 {
2164        return Err(FormulaDslError::InvalidArgument {
2165            reason: format!("{term_name}() requires degree >= 1: {raw}"),
2166        }
2167        .into());
2168    }
2169    let num_internal_knots =
2170        option_usize_strict(options, "internal_knots")?.unwrap_or(defaults.num_internal_knots);
2171    if num_internal_knots == 0 {
2172        return Err(FormulaDslError::InvalidArgument {
2173            reason: format!("{term_name}() requires internal_knots > 0: {raw}"),
2174        }
2175        .into());
2176    }
2177    // A DEFAULTED ladder adapts to the degree; an EXPLICIT one never does.
2178    //
2179    // `parse_linkwiggle_penalty_orders` falls back to
2180    // `WigglePenaltyConfig::cubic_triple_operator_default()` -- orders [1, 2, 3],
2181    // chosen for a CUBIC basis. A user writing `linkwiggle(degree=2, ...)`
2182    // overrides the degree and not the ladder, so the default asks for an
2183    // order-3 roughness penalty on a degree-2 basis. That penalty is the zero
2184    // matrix (the 3rd derivative of a degree-2 piecewise polynomial vanishes),
2185    // `bspline_unit_energy_factor` rightly refuses to build it, and the whole
2186    // fit dies with "Spline degree 2 is too low for derivative order 3".
2187    //
2188    // The user never asked for order 3 here, so dropping it is resolving a
2189    // default rather than overriding a request. When the ladder IS explicit the
2190    // orders stay exactly as written and the downstream refusal stands --
2191    // `wiggle::tests::unsupported_derivative_order_is_rejected_not_clamped`
2192    // pins that, and it is the reason this clip lives here at the
2193    // default-resolution seam instead of inside
2194    // `canonical_wiggle_function_penalties`, where it would have silently
2195    // clamped explicit requests too.
2196    let penalty_order_requested = options
2197        .get("penalty_order")
2198        .is_some_and(|raw| !raw.trim().is_empty());
2199    let mut penalty_orders =
2200        parse_linkwiggle_penalty_orders(options.get("penalty_order").map(String::as_str))?;
2201    if !penalty_order_requested {
2202        let supported: Vec<usize> = penalty_orders
2203            .iter()
2204            .copied()
2205            .filter(|&order| order <= degree)
2206            .collect();
2207        if supported.is_empty() {
2208            return Err(FormulaDslError::InvalidArgument {
2209                reason: format!(
2210                    "{term_name}(degree={degree}) supports no default penalty order: the default                      ladder {penalty_orders:?} contains no derivative order a degree-{degree}                      basis can penalize"
2211                ),
2212            }
2213            .into());
2214        }
2215        penalty_orders = supported;
2216    }
2217    let double_penalty =
2218        option_bool_strict(options, "double_penalty")?.unwrap_or(defaults.double_penalty);
2219    Ok(LinkWiggleFormulaSpec {
2220        degree,
2221        num_internal_knots,
2222        penalty_orders,
2223        double_penalty,
2224    })
2225}
2226
2227fn parse_link_formulaspec(
2228    options: &BTreeMap<String, String>,
2229    raw: &str,
2230) -> Result<LinkFormulaSpec, String> {
2231    let link = options
2232        .get("type")
2233        .map(|s| s.trim().to_string())
2234        .ok_or_else(|| FormulaDslError::MalformedConfig {
2235            reason: format!("link() requires type=<link-name>: {raw}"),
2236        })?;
2237    if link.is_empty() {
2238        return Err(FormulaDslError::MalformedConfig {
2239            reason: format!("link() requires a non-empty type: {raw}"),
2240        }
2241        .into());
2242    }
2243    let mixture_rho = options.get("rho").map(|s| s.trim().to_string());
2244    let sas_init = options.get("sas_init").map(|s| s.trim().to_string());
2245    let beta_logistic_init = options
2246        .get("beta_logistic_init")
2247        .map(|s| s.trim().to_string());
2248    Ok(LinkFormulaSpec {
2249        link,
2250        mixture_rho,
2251        sas_init,
2252        beta_logistic_init,
2253    })
2254}
2255
2256fn parse_survival_formulaspec(
2257    options: &BTreeMap<String, String>,
2258    raw: &str,
2259) -> Result<SurvivalFormulaSpec, String> {
2260    if options.is_empty() {
2261        return Err(FormulaDslError::MalformedConfig {
2262            reason: format!(
2263                "survmodel() requires at least one named option (e.g., spec=..., distribution=...): {raw}"
2264            ),
2265        }
2266        .into());
2267    }
2268    Ok(SurvivalFormulaSpec {
2269        spec: options.get("spec").map(|s| s.trim().to_string()),
2270        survival_distribution: options.get("distribution").map(|s| s.trim().to_string()),
2271    })
2272}
2273
2274fn parse_bounded_priorspec(
2275    options: &BTreeMap<String, String>,
2276    min: f64,
2277    max: f64,
2278    raw: &str,
2279) -> Result<BoundedCoefficientPriorSpec, String> {
2280    let prior_mode = options.get("prior").map(|s| s.to_ascii_lowercase());
2281    let pull = options.get("pull").map(|s| s.to_ascii_lowercase());
2282    let target = parse_optional_f64_option(options, "target", raw)?;
2283    let strength = parse_optional_f64_option(options, "strength", raw)?;
2284
2285    let target_mode = target.is_some() || strength.is_some();
2286    if prior_mode.is_some() && pull.is_some() {
2287        return Err(FormulaDslError::IncompatibleTerm {
2288            reason: format!("bounded() cannot combine prior=... with pull=...: {raw}"),
2289        }
2290        .into());
2291    }
2292    if prior_mode.is_some() && target_mode {
2293        return Err(FormulaDslError::IncompatibleTerm {
2294            reason: format!("bounded() cannot combine prior=... with target/strength: {raw}"),
2295        }
2296        .into());
2297    }
2298    if pull.is_some() && target_mode {
2299        return Err(FormulaDslError::IncompatibleTerm {
2300            reason: format!("bounded() cannot combine pull=... with target/strength: {raw}"),
2301        }
2302        .into());
2303    }
2304
2305    if let Some(priorname) = prior_mode {
2306        return match priorname.as_str() {
2307            "none" => Ok(BoundedCoefficientPriorSpec::None),
2308            "uniform" | "log-jacobian" | "log_jacobian" | "jacobian" => {
2309                Ok(BoundedCoefficientPriorSpec::Uniform)
2310            }
2311            "center" => Ok(BoundedCoefficientPriorSpec::Beta { a: 2.0, b: 2.0 }),
2312            _ => Err(FormulaDslError::InvalidArgument {
2313                reason: format!(
2314                    "bounded() prior must currently be one of none|uniform|log-jacobian|center, got '{}': {raw}",
2315                    priorname
2316                ),
2317            }
2318            .into()),
2319        };
2320    }
2321
2322    if let Some(pull_mode) = pull {
2323        return match pull_mode.as_str() {
2324            "uniform" | "log-jacobian" | "log_jacobian" | "jacobian" => {
2325                Ok(BoundedCoefficientPriorSpec::Uniform)
2326            }
2327            "center" => Ok(BoundedCoefficientPriorSpec::Beta { a: 2.0, b: 2.0 }),
2328            _ => Err(FormulaDslError::InvalidArgument {
2329                reason: format!(
2330                    "bounded() pull must currently be 'uniform'/'log-jacobian' or 'center', got '{}': {raw}",
2331                    pull_mode
2332                ),
2333            }
2334            .into()),
2335        };
2336    }
2337
2338    if target_mode {
2339        let targetvalue = target.ok_or_else(|| FormulaDslError::MalformedConfig {
2340            reason: format!("bounded() target is required with strength: {raw}"),
2341        })?;
2342        let strengthvalue = strength.ok_or_else(|| FormulaDslError::MalformedConfig {
2343            reason: format!("bounded() strength is required with target: {raw}"),
2344        })?;
2345        if !(min < targetvalue && targetvalue < max) {
2346            return Err(FormulaDslError::InvalidArgument {
2347                reason: format!("bounded() target must lie strictly inside ({min}, {max}): {raw}"),
2348            }
2349            .into());
2350        }
2351        if !strengthvalue.is_finite() || strengthvalue <= 0.0 {
2352            return Err(FormulaDslError::InvalidArgument {
2353                reason: format!("bounded() strength must be finite and > 0: {raw}"),
2354            }
2355            .into());
2356        }
2357        let z = (targetvalue - min) / (max - min);
2358        let a = 1.0 + strengthvalue * z;
2359        let b = 1.0 + strengthvalue * (1.0 - z);
2360        return Ok(BoundedCoefficientPriorSpec::Beta { a, b });
2361    }
2362
2363    Ok(BoundedCoefficientPriorSpec::None)
2364}
2365
2366// ---------------------------------------------------------------------------
2367// Top-level formula and term parsers
2368// ---------------------------------------------------------------------------
2369
2370pub fn formula_rhs_text(formula: &str) -> Result<String, String> {
2371    let parsed = parse_formula_dsl(formula)?;
2372    if parsed.rhs_terms.is_empty() {
2373        return Err(FormulaDslError::ParseError {
2374            reason: "formula right-hand side cannot be empty".to_string(),
2375        }
2376        .into());
2377    }
2378    Ok(parsed.rhs_terms.join(" + "))
2379}
2380
2381/// Parsed Surv(...) response specification.
2382///
2383/// `entry` is `None` for the 2-arg right-censored shorthand
2384/// `Surv(time, event)`, which matches the R survival/mgcv default: every
2385/// subject has entry time zero. Callers materialize a zero entry column
2386/// when this is `None`.
2387pub fn parse_surv_response(
2388    lhs: &str,
2389) -> Result<Option<(Option<String>, String, String)>, FormulaDslError> {
2390    let trimmed = lhs.trim();
2391    let call = match parse_function_call(trimmed) {
2392        Ok(call) => call,
2393        Err(_) => return Ok(None),
2394    };
2395    if !call.name.eq_ignore_ascii_case("surv") {
2396        return Ok(None);
2397    }
2398    let vars = call
2399        .args
2400        .iter()
2401        .filter_map(|arg| match arg {
2402            CallArgSpec::Positional(v) => Some(v.trim().to_string()),
2403            CallArgSpec::Named { .. } => None,
2404        })
2405        .filter(|s| !s.is_empty())
2406        .collect::<Vec<_>>();
2407    match vars.len() {
2408        // Right-censored shorthand: Surv(time, event) ≡ Surv(0, time, event)
2409        // with a synthetic zero entry column. This matches R's
2410        // `survival::Surv(time, event)` default for left-truncation-free data.
2411        2 => Ok(Some((None, vars[0].clone(), vars[1].clone()))),
2412        3 => Ok(Some((
2413            Some(vars[0].clone()),
2414            vars[1].clone(),
2415            vars[2].clone(),
2416        ))),
2417        n => Err(FormulaDslError::InvalidArgument {
2418            reason: format!(
2419                "Surv(...) expects either Surv(time, event) (right-censored) or \
2420                 Surv(entry, exit, event) (left-truncated); got {n} columns"
2421            ),
2422        }),
2423    }
2424}
2425
2426/// Parsed `SurvInterval(L, R, event)` interval-censored response.
2427///
2428/// Returns `Some((left_col, right_col, event_col))` when the left-hand side is a
2429/// `SurvInterval(...)` call, `None` otherwise (so a plain `Surv(...)` or a bare
2430/// column response falls through to the other response parsers).
2431///
2432/// Interval censoring observes only a bracket `T ∈ (L, R]` — the exact event
2433/// time is never seen — and its row contribution is the survival-mass difference
2434/// `log[S(L) − S(R)]`, distinct from both the exact-event point density and the
2435/// single-sided right-censored survival. A *dedicated call name* (rather than
2436/// overloading the 3-argument `Surv(entry, exit, event)` delayed-entry form,
2437/// which is also 3-argument and semantically incompatible) is the unambiguous
2438/// DSL spelling: it mirrors flexsurv's `Surv(L, R, type="interval2")` intent
2439/// without colliding with the existing left-truncation grammar.
2440pub fn parse_surv_interval_response(
2441    lhs: &str,
2442) -> Result<Option<(String, String, String)>, FormulaDslError> {
2443    let trimmed = lhs.trim();
2444    let call = match parse_function_call(trimmed) {
2445        Ok(call) => call,
2446        Err(_) => return Ok(None),
2447    };
2448    if !call.name.eq_ignore_ascii_case("survinterval") {
2449        return Ok(None);
2450    }
2451    let vars = call
2452        .args
2453        .iter()
2454        .filter_map(|arg| match arg {
2455            CallArgSpec::Positional(v) => Some(v.trim().to_string()),
2456            CallArgSpec::Named { .. } => None,
2457        })
2458        .filter(|s| !s.is_empty())
2459        .collect::<Vec<_>>();
2460    match vars.len() {
2461        3 => Ok(Some((vars[0].clone(), vars[1].clone(), vars[2].clone()))),
2462        n => Err(FormulaDslError::InvalidArgument {
2463            reason: format!(
2464                "SurvInterval(...) expects SurvInterval(L, R, event) (interval-censored, the \
2465                 observed bracket T ∈ (L, R]); got {n} columns"
2466            ),
2467        }),
2468    }
2469}
2470
2471fn top_level_formula_separator(input: &str) -> Result<Option<usize>, String> {
2472    let mut depth = 0_i32;
2473    let mut in_single = false;
2474    let mut in_double = false;
2475
2476    for (idx, ch) in input.char_indices() {
2477        let quoted = in_single || in_double;
2478        if ch == '\'' && !in_double {
2479            in_single = !in_single;
2480        } else if ch == '"' && !in_single {
2481            in_double = !in_double;
2482        } else if !quoted && matches!(ch, '(' | '[' | '{') {
2483            depth += 1;
2484        } else if !quoted && matches!(ch, ')' | ']' | '}') && depth > 0 {
2485            depth -= 1;
2486        } else if ch == '~' && !quoted && depth == 0 {
2487            return Ok(Some(idx));
2488        }
2489    }
2490
2491    if in_single || in_double || depth != 0 {
2492        return Err(FormulaDslError::ParseError {
2493            reason: "invalid auxiliary formula syntax: unbalanced parentheses or quotes"
2494                .to_string(),
2495        }
2496        .into());
2497    }
2498    Ok(None)
2499}
2500
2501pub fn parse_matching_auxiliary_formula(
2502    formula: &str,
2503    response: &str,
2504    flag_name: &str,
2505) -> Result<(String, ParsedFormula), FormulaDslError> {
2506    let rhs = formula.trim();
2507    if top_level_formula_separator(rhs)?.is_some() {
2508        return Err(FormulaDslError::InvalidArgument {
2509            reason: format!(
2510                "{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)"
2511            ),
2512        });
2513    }
2514    let parsed_formula = parse_formula(&format!("{response} ~ {rhs}"))?;
2515    Ok((rhs.to_string(), parsed_formula))
2516}
2517
2518pub fn validate_auxiliary_formula_controls(
2519    parsed_formula: &ParsedFormula,
2520    flag_name: &str,
2521) -> Result<(), String> {
2522    if parsed_formula.linkwiggle.is_some() {
2523        return Err(FormulaDslError::IncompatibleTerm {
2524            reason: format!(
2525                "linkwiggle(...) is only supported in the main formula, not {flag_name}"
2526            ),
2527        }
2528        .into());
2529    }
2530    if parsed_formula.timewiggle.is_some() {
2531        return Err(FormulaDslError::IncompatibleTerm {
2532            reason: format!(
2533                "timewiggle(...) is only supported in the main survival formula, not {flag_name}"
2534            ),
2535        }
2536        .into());
2537    }
2538    if parsed_formula.linkspec.is_some() {
2539        return Err(FormulaDslError::IncompatibleTerm {
2540            reason: format!("link(...) is only supported in the main formula, not {flag_name}"),
2541        }
2542        .into());
2543    }
2544    if parsed_formula.survivalspec.is_some() {
2545        return Err(FormulaDslError::IncompatibleTerm {
2546            reason: format!(
2547                "survmodel(...) is only supported in the main survival formula, not {flag_name}"
2548            ),
2549        }
2550        .into());
2551    }
2552    if !parsed_formula.logslope_surfaces.is_empty() && flag_name != "--logslope-formula" {
2553        return Err(FormulaDslError::IncompatibleTerm {
2554            reason: format!(
2555                "logslope(...) is only supported in --logslope-formula, not {flag_name}"
2556            ),
2557        }
2558        .into());
2559    }
2560    Ok::<(), _>(())
2561}
2562
2563pub fn parse_formula(formula: &str) -> Result<ParsedFormula, FormulaDslError> {
2564    let parsed_dsl =
2565        parse_formula_dsl(formula).map_err(|reason| FormulaDslError::ParseError { reason })?;
2566    let lhs = parsed_dsl.response_expr.trim();
2567    if lhs.is_empty() {
2568        return Err(FormulaDslError::ParseError {
2569            reason: "formula response (left-hand side) cannot be empty".to_string(),
2570        });
2571    }
2572    let mut terms = Vec::<ParsedTerm>::new();
2573    let mut linkwiggle: Option<LinkWiggleFormulaSpec> = None;
2574    let mut timewiggle: Option<LinkWiggleFormulaSpec> = None;
2575    let mut linkspec: Option<LinkFormulaSpec> = None;
2576    let mut survivalspec: Option<SurvivalFormulaSpec> = None;
2577    let mut logslope_surfaces = Vec::<LogSlopeSurfaceSpec>::new();
2578    // Track seen-term-keys so we can reject exact duplicates like
2579    // `y ~ smooth(x) + smooth(x)` upfront — without this the duplicate
2580    // produces a rank-deficient design and the user has no idea why their
2581    // fit is over-parameterized.
2582    let mut seen_term_keys: std::collections::BTreeSet<String> = std::collections::BTreeSet::new();
2583    let mut expanded_terms = Vec::<String>::new();
2584    for raw in parsed_dsl.rhs_terms {
2585        let trimmed = raw.trim();
2586        if trimmed.is_empty() {
2587            expanded_terms.push(String::new());
2588            continue;
2589        }
2590        // Single function-call terms (smooths, group(), etc.) are opaque to
2591        // WR expansion; pass them through verbatim so parse_term sees the
2592        // exact source string. Bare identifiers and any term mentioning
2593        // `:`, `*`, `/`, `^` are routed through the AST-driven WR expander.
2594        let is_call = parse_function_call(trimmed).is_ok();
2595        let needs_expansion = !is_call
2596            && trimmed
2597                .chars()
2598                .scan(0i32, |depth, ch| {
2599                    let d_before = *depth;
2600                    if matches!(ch, '(' | '[' | '{') {
2601                        *depth += 1;
2602                    } else if matches!(ch, ')' | ']' | '}') && *depth > 0 {
2603                        *depth -= 1;
2604                    }
2605                    Some((d_before, ch))
2606                })
2607                .any(|(d, ch)| d == 0 && matches!(ch, ':' | '*' | '/' | '^'));
2608        if needs_expansion {
2609            for atoms in
2610                expand_wr_term(trimmed).map_err(|reason| FormulaDslError::ParseError { reason })?
2611            {
2612                if atoms.is_empty() {
2613                    continue;
2614                }
2615                expanded_terms.push(atoms.join(":"));
2616            }
2617        } else {
2618            expanded_terms.push(trimmed.to_string());
2619        }
2620    }
2621
2622    for raw in expanded_terms {
2623        let t = raw.trim();
2624        if t.is_empty() || t == "1" {
2625            continue;
2626        }
2627        if t == "0" || t == "-1" {
2628            return Err(FormulaDslError::IncompatibleTerm {
2629                reason: "formula terms '0'/'-1' (intercept removal) are not supported yet"
2630                    .to_string(),
2631            });
2632        }
2633        // Normalize whitespace so `smooth(x)` and `smooth( x )` match,
2634        // but preserve whitespace inside string literals so that
2635        // `bs="a b"` and `bs="ab"` do not collide.
2636        let key: String = {
2637            let mut acc = String::with_capacity(t.len());
2638            let mut in_single = false;
2639            let mut in_double = false;
2640            for ch in t.chars() {
2641                match ch {
2642                    '\'' if !in_double => {
2643                        in_single = !in_single;
2644                        acc.push(ch);
2645                    }
2646                    '"' if !in_single => {
2647                        in_double = !in_double;
2648                        acc.push(ch);
2649                    }
2650                    c if c.is_whitespace() && !in_single && !in_double => {}
2651                    _ => acc.push(ch),
2652                }
2653            }
2654            acc
2655        };
2656        if !seen_term_keys.insert(key.clone()) {
2657            return Err(FormulaDslError::IncompatibleTerm {
2658                reason: format!(
2659                    "formula `{formula}` lists term `{t}` more than once. \
2660                     Duplicate terms produce a rank-deficient design; \
2661                     keep one copy or differentiate them (e.g. distinct k=, bs= options)."
2662                ),
2663            });
2664        }
2665        match parse_term(t)? {
2666            ParsedTerm::LinkWiggle { options } => {
2667                if linkwiggle.is_some() {
2668                    return Err(FormulaDslError::IncompatibleTerm {
2669                        reason: "formula can include at most one linkwiggle(...) term".to_string(),
2670                    });
2671                }
2672                linkwiggle = Some(parse_linkwiggle_formulaspec(&options, t)?);
2673            }
2674            ParsedTerm::TimeWiggle { options } => {
2675                if timewiggle.is_some() {
2676                    return Err(FormulaDslError::IncompatibleTerm {
2677                        reason: "formula can include at most one timewiggle(...) term".to_string(),
2678                    });
2679                }
2680                timewiggle = Some(parse_linkwiggle_formulaspec(&options, t)?);
2681            }
2682            ParsedTerm::LinkConfig { options } => {
2683                if linkspec.is_some() {
2684                    return Err(FormulaDslError::IncompatibleTerm {
2685                        reason: "formula can include at most one link(...) term".to_string(),
2686                    });
2687                }
2688                linkspec = Some(parse_link_formulaspec(&options, t)?);
2689            }
2690            ParsedTerm::SurvivalConfig { options } => {
2691                if survivalspec.is_some() {
2692                    return Err(FormulaDslError::IncompatibleTerm {
2693                        reason: "formula can include at most one survmodel(...) term".to_string(),
2694                    });
2695                }
2696                survivalspec = Some(parse_survival_formulaspec(&options, t)?);
2697            }
2698            ParsedTerm::LogSlopeSurface { z_column, terms } => {
2699                logslope_surfaces.push(LogSlopeSurfaceSpec { z_column, terms });
2700            }
2701            other => terms.push(other),
2702        }
2703    }
2704    // Reject self-referential formulas like `y ~ smooth(y)` or `y ~ y`: the
2705    // response is its own predictor, which is a trivial identity fit and
2706    // almost certainly a user mistake. Only flag the simple-identifier case
2707    // (so Surv(entry, exit, event) ~ smooth(time) is left alone — the
2708    // response is the Surv triple, not the bare "time" column).
2709    if lhs.chars().all(|c| c.is_alphanumeric() || c == '_')
2710        && !lhs.is_empty()
2711        && parsed_terms_reference_column(&terms, lhs)
2712    {
2713        return Err(FormulaDslError::IncompatibleTerm {
2714            reason: format!(
2715                "formula `{formula}` uses response column `{lhs}` as its own predictor. \
2716                 This fits y as a function of itself and is almost certainly a typo. \
2717                 Drop the term that mentions `{lhs}` from the right-hand side."
2718            ),
2719        });
2720    }
2721    Ok(ParsedFormula {
2722        response: lhs.to_string(),
2723        terms,
2724        logslope_surfaces,
2725        linkwiggle,
2726        timewiggle,
2727        linkspec,
2728        survivalspec,
2729    })
2730}
2731
2732pub fn parse_term(raw: &str) -> Result<ParsedTerm, String> {
2733    fn split_call_args(call: &FunctionCallSpec) -> (Vec<String>, BTreeMap<String, String>) {
2734        let mut vars = Vec::<String>::new();
2735        let mut options = BTreeMap::<String, String>::new();
2736        for arg in &call.args {
2737            match arg {
2738                CallArgSpec::Positional(v) => vars.push(v.trim().to_string()),
2739                CallArgSpec::Named { key, value } => {
2740                    options.insert(key.to_ascii_lowercase(), strip_quotes(value).to_string());
2741                }
2742            }
2743        }
2744        (vars, options)
2745    }
2746
2747    // Wilkinson-Rogers `:` interaction term. The expander in `parse_formula`
2748    // produces `a:b[:c...]` for these; parse_term is also reached directly
2749    // from tests, so handle the syntax here as well.
2750    if raw.contains(':')
2751        && !raw.contains('(')
2752        && raw.split(':').all(|piece| is_exact_ident(piece.trim()))
2753    {
2754        let vars: Vec<String> = raw
2755            .split(':')
2756            .map(|piece| piece.trim().to_string())
2757            .collect();
2758        if vars.len() >= 2 {
2759            let mut sorted = vars.clone();
2760            sorted.sort();
2761            sorted.dedup();
2762            if sorted.len() != vars.len() {
2763                return Err(FormulaDslError::IncompatibleTerm {
2764                    reason: format!(
2765                        "interaction term `{raw}` references the same variable more than once"
2766                    ),
2767                }
2768                .into());
2769            }
2770            return Ok(ParsedTerm::Interaction {
2771                vars: sorted,
2772                double_penalty: false,
2773            });
2774        }
2775    }
2776
2777    // A term that does not parse as `name(args)` is not a function-call term;
2778    // the plain-variable handling below is the answer, so there is no error here
2779    // to report.
2780    if let Ok(call) = parse_function_call(raw) {
2781        let name = call.name.to_ascii_lowercase();
2782        let (vars, mut options) = split_call_args(&call);
2783        match name.as_str() {
2784            "constrain" | "constraint" | "box" => {
2785                if vars.len() != 1 {
2786                    return Err(FormulaDslError::InvalidArgument {
2787                        reason: format!(
2788                            "constrain()/constraint()/box() expects exactly one variable: {raw}"
2789                        ),
2790                    }
2791                    .into());
2792                }
2793                validate_known_term_options(
2794                    "constrain",
2795                    &options,
2796                    &["min", "lower", "max", "upper", "double_penalty"],
2797                    raw,
2798                )?;
2799                let (coefficient_min, coefficient_max) =
2800                    parse_linear_constraint_bounds(&options, raw)?;
2801                if coefficient_min.is_none() && coefficient_max.is_none() {
2802                    return Err(FormulaDslError::MalformedConfig {
2803                        reason: format!(
2804                            "constrain()/constraint()/box() requires at least one of min/lower/max/upper: {raw}"
2805                        ),
2806                    }
2807                    .into());
2808                }
2809                return Ok(ParsedTerm::Linear {
2810                    name: vars[0].clone(),
2811                    explicit: true,
2812                    double_penalty: option_bool_strict(&options, "double_penalty")?
2813                        .unwrap_or(false),
2814                    coefficient_min,
2815                    coefficient_max,
2816                });
2817            }
2818            "nonnegative" | "nonnegative_coef" => {
2819                if vars.len() != 1 {
2820                    return Err(FormulaDslError::InvalidArgument {
2821                        reason: format!("nonnegative() expects exactly one variable: {raw}"),
2822                    }
2823                    .into());
2824                }
2825                validate_known_term_options("nonnegative", &options, &["double_penalty"], raw)?;
2826                return Ok(ParsedTerm::Linear {
2827                    name: vars[0].clone(),
2828                    explicit: true,
2829                    double_penalty: option_bool_strict(&options, "double_penalty")?
2830                        .unwrap_or(false),
2831                    coefficient_min: Some(0.0),
2832                    coefficient_max: None,
2833                });
2834            }
2835            "nonpositive" | "nonpositive_coef" => {
2836                if vars.len() != 1 {
2837                    return Err(FormulaDslError::InvalidArgument {
2838                        reason: format!("nonpositive() expects exactly one variable: {raw}"),
2839                    }
2840                    .into());
2841                }
2842                validate_known_term_options("nonpositive", &options, &["double_penalty"], raw)?;
2843                return Ok(ParsedTerm::Linear {
2844                    name: vars[0].clone(),
2845                    explicit: true,
2846                    double_penalty: option_bool_strict(&options, "double_penalty")?
2847                        .unwrap_or(false),
2848                    coefficient_min: None,
2849                    coefficient_max: Some(0.0),
2850                });
2851            }
2852            "bounded" => {
2853                if vars.len() != 1 {
2854                    return Err(FormulaDslError::InvalidArgument {
2855                        reason: format!("bounded() expects exactly one variable: {raw}"),
2856                    }
2857                    .into());
2858                }
2859                validate_known_term_options(
2860                    "bounded",
2861                    &options,
2862                    &[
2863                        "min",
2864                        "max",
2865                        "prior",
2866                        "pull",
2867                        "target",
2868                        "strength",
2869                        "double_penalty",
2870                    ],
2871                    raw,
2872                )?;
2873                let min = parse_required_f64_option(&options, "min", raw)?;
2874                let max = parse_required_f64_option(&options, "max", raw)?;
2875                if !min.is_finite() || !max.is_finite() || min >= max {
2876                    return Err(FormulaDslError::InvalidArgument {
2877                        reason: format!(
2878                            "bounded() requires finite min < max, got min={min}, max={max}: {raw}"
2879                        ),
2880                    }
2881                    .into());
2882                }
2883                let prior = parse_bounded_priorspec(&options, min, max, raw)?;
2884                return Ok(ParsedTerm::BoundedLinear {
2885                    name: vars[0].clone(),
2886                    min,
2887                    max,
2888                    prior,
2889                    // Unlike a plain `linear()` term, `bounded()` already commits
2890                    // the coefficient to an exact interval transform (plus an
2891                    // optional prior); layering the null-space ridge on top is
2892                    // structurally rejected downstream (`design_construction.rs`:
2893                    // "bounded linear term ... cannot also use double_penalty"),
2894                    // so the default must be `false`, not the `linear()`/`s()`
2895                    // convention of `true`.
2896                    double_penalty: option_bool_strict(&options, "double_penalty")?
2897                        .unwrap_or(false),
2898                });
2899            }
2900            "group" | "re" | "factor" => {
2901                if vars.len() != 1 {
2902                    return Err(FormulaDslError::InvalidArgument {
2903                        reason: format!(
2904                            "{name}() expects exactly one variable, got '{}': {raw}",
2905                            vars.join(",")
2906                        ),
2907                    }
2908                    .into());
2909                }
2910                // `factor(g)` is a FIXED categorical factor (R `factor()` /
2911                // patsy `C()`): it forces categorical encoding of the column
2912                // but, like a bare `+ g` main effect, is strict on unseen
2913                // levels. `group(g)`/`re(g)` are genuine random effects that
2914                // shrink a held-out group to the population mean, so they
2915                // tolerate unseen levels. Both share the penalized-categorical
2916                // block; only the unseen policy differs (#2137/#2102).
2917                let lenient_unseen = name != "factor";
2918                return Ok(ParsedTerm::RandomEffect {
2919                    name: vars[0].clone(),
2920                    lenient_unseen,
2921                });
2922            }
2923            "tensor" | "interaction" | "te" => {
2924                if vars.len() < 2 {
2925                    return Err(FormulaDslError::InvalidArgument {
2926                        reason: format!(
2927                            "tensor()/interaction()/te() requires at least two variables: {raw}"
2928                        ),
2929                    }
2930                    .into());
2931                }
2932                return Ok(ParsedTerm::Smooth {
2933                    label: raw.to_string(),
2934                    vars,
2935                    kind: SmoothKind::Te,
2936                    options,
2937                });
2938            }
2939            "t2" => {
2940                if vars.len() < 2 {
2941                    return Err(FormulaDslError::InvalidArgument {
2942                        reason: format!("t2() requires at least two variables: {raw}"),
2943                    }
2944                    .into());
2945                }
2946                return Ok(ParsedTerm::Smooth {
2947                    label: raw.to_string(),
2948                    vars,
2949                    kind: SmoothKind::T2,
2950                    options,
2951                });
2952            }
2953            "ti" => {
2954                // Tensor interaction smooth (mgcv `ti`): structurally a
2955                // tensor-product smooth, but the marginal main effects are
2956                // excluded so only the pure interaction is modeled. Shares the
2957                // tensor materialization path with `te`; the distinct
2958                // `SmoothKind::Ti` drives per-margin sum-to-zero
2959                // identifiability in the term builder.
2960                if vars.len() < 2 {
2961                    return Err(FormulaDslError::InvalidArgument {
2962                        reason: format!("ti() requires at least two variables: {raw}"),
2963                    }
2964                    .into());
2965                }
2966                return Ok(ParsedTerm::Smooth {
2967                    label: raw.to_string(),
2968                    vars,
2969                    kind: SmoothKind::Ti,
2970                    options,
2971                });
2972            }
2973            "fs" | "sz" => {
2974                if vars.len() != 2 {
2975                    return Err(format!("{}() expects exactly two variables: {raw}", name));
2976                }
2977                options.insert("bs".to_string(), name.clone());
2978                return Ok(ParsedTerm::Smooth {
2979                    label: raw.to_string(),
2980                    vars,
2981                    kind: SmoothKind::S,
2982                    options,
2983                });
2984            }
2985            "thinplate" | "thin_plate" | "tps" => {
2986                if vars.len() < 2 {
2987                    return Err(FormulaDslError::InvalidArgument {
2988                        reason: format!(
2989                            "thinplate()/thin_plate()/tps() requires at least two variables: {raw}"
2990                        ),
2991                    }
2992                    .into());
2993                }
2994                options.insert("type".to_string(), "tps".to_string());
2995                return Ok(ParsedTerm::Smooth {
2996                    label: raw.to_string(),
2997                    vars,
2998                    kind: SmoothKind::S,
2999                    options,
3000                });
3001            }
3002            "smooth" | "s" | "cyclic" | "periodic" | "cc" | "cp" => {
3003                if vars.is_empty() {
3004                    return Err(FormulaDslError::InvalidArgument {
3005                        reason: format!("smooth()/s() requires at least one variable: {raw}"),
3006                    }
3007                    .into());
3008                }
3009                // mgcv idiom: `s(g, bs='re')` with a single variable is a
3010                // random intercept on the factor `g`. Route it to the
3011                // dedicated random-effect machinery (which expects a single
3012                // categorical column) rather than to the factor-smooth path
3013                // (which requires a numeric companion).
3014                let bs_is_re = options
3015                    .get("bs")
3016                    .or_else(|| options.get("type"))
3017                    .map(|v| {
3018                        v.trim()
3019                            .trim_matches(|c| c == '\'' || c == '"')
3020                            .to_ascii_lowercase()
3021                    })
3022                    .as_deref()
3023                    == Some("re");
3024                if bs_is_re && vars.len() == 1 {
3025                    // `s(g, bs="re")` is a genuine random effect: lenient on
3026                    // unseen levels (held-out group → population mean).
3027                    return Ok(ParsedTerm::RandomEffect {
3028                        name: vars[0].clone(),
3029                        lenient_unseen: true,
3030                    });
3031                }
3032                if matches!(name.as_str(), "cyclic" | "periodic" | "cc" | "cp") {
3033                    options.insert("type".to_string(), "cyclic".to_string());
3034                }
3035                if matches!(name.as_str(), "fs" | "sz") {
3036                    options.insert("bs".to_string(), name.clone());
3037                }
3038                return Ok(ParsedTerm::Smooth {
3039                    label: raw.to_string(),
3040                    vars,
3041                    kind: SmoothKind::S,
3042                    options,
3043                });
3044            }
3045            "sphere" | "sos" | "spherical" | "s2" => {
3046                // `s2()` is an alias for the intrinsic S² (sphere) smooth, just
3047                // like `sphere()`/`sos()`/`spherical()`. All four share the
3048                // Wahba/harmonic sphere basis, so they must dispatch through
3049                // the identical `type=sphere` route; otherwise `s2()` would
3050                // silently fall back to a generic Euclidean 2-D smooth over
3051                // (lat, lon) and diverge in the spatial-kappa optimizer.
3052                if vars.len() != 2 {
3053                    return Err(FormulaDslError::InvalidArgument {
3054                        reason: format!(
3055                            "{name}() expects exactly two variables: latitude and longitude; got {} in {raw}",
3056                            vars.len()
3057                        ),
3058                    }
3059                    .into());
3060                }
3061                options.insert("type".to_string(), "sphere".to_string());
3062                return Ok(ParsedTerm::Smooth {
3063                    label: raw.to_string(),
3064                    vars,
3065                    kind: SmoothKind::S,
3066                    options,
3067                });
3068            }
3069            "mjs" | "measurejet" | "measure_jet" | "web" => {
3070                // Measure-jet spline smooth (`basis::measure_jet_smooth` docs)
3071                // for responses varying along an unknown low-dimensional set
3072                // inside a higher-dimensional ambient space. All aliases
3073                // dispatch through the identical `type=measurejet` route,
3074                // mirroring the sphere/curvature alias rule.
3075                if vars.is_empty() {
3076                    return Err(FormulaDslError::InvalidArgument {
3077                        reason: format!("{name}() requires at least one variable: {raw}"),
3078                    }
3079                    .into());
3080                }
3081                options.insert("type".to_string(), "measurejet".to_string());
3082                return Ok(ParsedTerm::Smooth {
3083                    label: raw.to_string(),
3084                    vars,
3085                    kind: SmoothKind::S,
3086                    options,
3087                });
3088            }
3089            "curv" | "curvature" | "constant_curvature" | "mkappa" => {
3090                // Constant-curvature (M_κ) geodesic-kernel smooth (#944): the
3091                // κ-generic sibling of sphere()/s2(), interpolating
3092                // S^d → ℝ^d → H^d through `kappa=` (default 0 = flat). All
3093                // four aliases must dispatch through the identical
3094                // `type=curvature` route, mirroring the sphere alias rule.
3095                if vars.is_empty() {
3096                    return Err(FormulaDslError::InvalidArgument {
3097                        reason: format!("{name}() requires at least one variable: {raw}"),
3098                    }
3099                    .into());
3100                }
3101                options.insert("type".to_string(), "curvature".to_string());
3102                return Ok(ParsedTerm::Smooth {
3103                    label: raw.to_string(),
3104                    vars,
3105                    kind: SmoothKind::S,
3106                    options,
3107                });
3108            }
3109            "matern" => {
3110                if vars.is_empty() {
3111                    return Err(FormulaDslError::InvalidArgument {
3112                        reason: format!("matern() requires at least one variable: {raw}"),
3113                    }
3114                    .into());
3115                }
3116                options.insert("type".to_string(), "matern".to_string());
3117                return Ok(ParsedTerm::Smooth {
3118                    label: raw.to_string(),
3119                    vars,
3120                    kind: SmoothKind::S,
3121                    options,
3122                });
3123            }
3124            "duchon" => {
3125                if vars.is_empty() {
3126                    return Err(FormulaDslError::InvalidArgument {
3127                        reason: format!("duchon() requires at least one variable: {raw}"),
3128                    }
3129                    .into());
3130                }
3131                if option_bool(&options, "cyclic").unwrap_or(false)
3132                    || option_bool(&options, "periodic").unwrap_or(false)
3133                {
3134                    options.insert("cyclic".to_string(), "true".to_string());
3135                }
3136                options.insert("type".to_string(), "duchon".to_string());
3137                return Ok(ParsedTerm::Smooth {
3138                    label: raw.to_string(),
3139                    vars,
3140                    kind: SmoothKind::S,
3141                    options,
3142                });
3143            }
3144            "pca" => {
3145                if vars.is_empty() {
3146                    return Err(FormulaDslError::InvalidArgument {
3147                        reason: format!("pca() requires at least one variable: {raw}"),
3148                    }
3149                    .into());
3150                }
3151                options.insert("type".to_string(), "pca".to_string());
3152                return Ok(ParsedTerm::Smooth {
3153                    label: raw.to_string(),
3154                    vars,
3155                    kind: SmoothKind::S,
3156                    options,
3157                });
3158            }
3159            "linkwiggle" => {
3160                if !vars.is_empty() {
3161                    return Err(FormulaDslError::InvalidArgument {
3162                        reason: format!(
3163                            "linkwiggle() takes named options only; positional args are not supported: {raw}"
3164                        ),
3165                    }
3166                    .into());
3167                }
3168                return Ok(ParsedTerm::LinkWiggle { options });
3169            }
3170            "timewiggle" => {
3171                if !vars.is_empty() {
3172                    return Err(FormulaDslError::InvalidArgument {
3173                        reason: format!(
3174                            "timewiggle() takes named options only; positional args are not supported: {raw}"
3175                        ),
3176                    }
3177                    .into());
3178                }
3179                return Ok(ParsedTerm::TimeWiggle { options });
3180            }
3181            "link" => {
3182                if !vars.is_empty() {
3183                    return Err(FormulaDslError::InvalidArgument {
3184                        reason: format!(
3185                            "link() takes named options only; positional args are not supported: {raw}"
3186                        ),
3187                    }
3188                    .into());
3189                }
3190                return Ok(ParsedTerm::LinkConfig { options });
3191            }
3192            "survmodel" => {
3193                if !vars.is_empty() {
3194                    return Err(FormulaDslError::InvalidArgument {
3195                        reason: format!(
3196                            "survmodel() takes named options only; positional args are not supported: {raw}"
3197                        ),
3198                    }
3199                    .into());
3200                }
3201                return Ok(ParsedTerm::SurvivalConfig { options });
3202            }
3203            "logslope" | "log_slope" | "log_slope_surface" => {
3204                validate_known_term_options("logslope", &options, &[], raw)?;
3205                if vars.len() < 2 {
3206                    return Err(FormulaDslError::InvalidArgument {
3207                        reason: format!(
3208                            "logslope() expects a z column followed by one or more RHS terms; add one logslope(z, ...) declaration per vector-z coordinate: {raw}"
3209                        ),
3210                    }
3211                    .into());
3212                }
3213                let z_column = vars[0].trim();
3214                if !is_exact_ident(z_column) {
3215                    return Err(FormulaDslError::InvalidArgument {
3216                        reason: format!(
3217                            "logslope() z column must be a bare column name, got `{z_column}` in {raw}"
3218                        ),
3219                    }
3220                    .into());
3221                }
3222                let rhs = vars[1..].join(" + ");
3223                let parsed = parse_formula(&format!("__logslope__ ~ {rhs}"))?;
3224                if !parsed.logslope_surfaces.is_empty() {
3225                    return Err(FormulaDslError::IncompatibleTerm {
3226                        reason: format!(
3227                            "logslope() declarations cannot be nested inside another logslope(): {raw}"
3228                        ),
3229                    }
3230                    .into());
3231                }
3232                validate_auxiliary_formula_controls(&parsed, "logslope()")?;
3233                return Ok(ParsedTerm::LogSlopeSurface {
3234                    z_column: z_column.to_string(),
3235                    terms: parsed.terms,
3236                });
3237            }
3238            "linear" => {
3239                if vars.len() != 1 {
3240                    return Err(FormulaDslError::InvalidArgument {
3241                        reason: format!("linear() expects exactly one variable: {raw}"),
3242                    }
3243                    .into());
3244                }
3245                validate_known_term_options(
3246                    "linear",
3247                    &options,
3248                    &["min", "lower", "max", "upper", "double_penalty"],
3249                    raw,
3250                )?;
3251                let (coefficient_min, coefficient_max) =
3252                    parse_linear_constraint_bounds(&options, raw)?;
3253                let double_penalty =
3254                    option_bool_strict(&options, "double_penalty")?.unwrap_or(false);
3255                if vars[0].contains(':') {
3256                    if coefficient_min.is_some() || coefficient_max.is_some() {
3257                        return Err(FormulaDslError::IncompatibleTerm {
3258                            reason: format!(
3259                                "linear() coefficient bounds are not supported on an interaction: {raw}"
3260                            ),
3261                        }
3262                        .into());
3263                    }
3264                    let mut interaction_vars = vars[0]
3265                        .split(':')
3266                        .map(str::trim)
3267                        .map(str::to_string)
3268                        .collect::<Vec<_>>();
3269                    if interaction_vars.len() < 2
3270                        || interaction_vars.iter().any(|var| !is_exact_ident(var))
3271                    {
3272                        return Err(FormulaDslError::InvalidArgument {
3273                            reason: format!(
3274                                "linear() interaction must contain at least two bare column names: {raw}"
3275                            ),
3276                        }
3277                        .into());
3278                    }
3279                    interaction_vars.sort();
3280                    let original_len = interaction_vars.len();
3281                    interaction_vars.dedup();
3282                    if interaction_vars.len() != original_len {
3283                        return Err(FormulaDslError::IncompatibleTerm {
3284                            reason: format!(
3285                                "linear() interaction references the same variable more than once: {raw}"
3286                            ),
3287                        }
3288                        .into());
3289                    }
3290                    return Ok(ParsedTerm::Interaction {
3291                        vars: interaction_vars,
3292                        double_penalty,
3293                    });
3294                }
3295                return Ok(ParsedTerm::Linear {
3296                    name: vars[0].clone(),
3297                    explicit: true,
3298                    double_penalty,
3299                    coefficient_min,
3300                    coefficient_max,
3301                });
3302            }
3303            _ => {
3304                return Err(format!(
3305                    "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()"
3306                ));
3307            }
3308        }
3309    }
3310
3311    let ident = raw.trim();
3312    if !is_exact_ident(ident) {
3313        return Err(FormulaDslError::UnknownIdentifier {
3314            reason: format!("unsupported top-level RHS term: {raw}"),
3315        }
3316        .into());
3317    }
3318
3319    Ok(ParsedTerm::Linear {
3320        name: ident.to_string(),
3321        explicit: false,
3322        double_penalty: false,
3323        coefficient_min: None,
3324        coefficient_max: None,
3325    })
3326}
3327
3328// ---------------------------------------------------------------------------
3329// Link choice parsing
3330// ---------------------------------------------------------------------------
3331
3332pub fn parse_link_choice(
3333    raw: Option<&str>,
3334    flexible_flag: bool,
3335) -> Result<Option<LinkChoice>, FormulaDslError> {
3336    if raw.is_none() && !flexible_flag {
3337        return Ok(None);
3338    }
3339    let Some(v) = raw else {
3340        return Ok(Some(LinkChoice {
3341            mode: LinkMode::Flexible,
3342            link: LinkFunction::Probit,
3343            mixture_components: None,
3344        }));
3345    };
3346    let t = v.trim().to_ascii_lowercase();
3347    if let Some(inner) = t
3348        .strip_prefix("flexible(")
3349        .and_then(|s| s.strip_suffix(')'))
3350    {
3351        if let Some(components_inner) = inner
3352            .strip_prefix("blended(")
3353            .and_then(|s| s.strip_suffix(')'))
3354            .or_else(|| {
3355                inner
3356                    .strip_prefix("mixture(")
3357                    .and_then(|s| s.strip_suffix(')'))
3358            })
3359        {
3360            parse_link_component_list(components_inner)?;
3361            return Err(FormulaDslError::IncompatibleTerm {
3362                reason:
3363                    "flexible(...) does not support blended(...)/mixture(...) links; wiggle is only supported for jointly fit standard links"
3364                        .to_string(),
3365            });
3366        }
3367        let link = parse_linkname(inner)?;
3368        if !linkname_supports_joint_wiggle(link) {
3369            return Err(FormulaDslError::IncompatibleTerm {
3370                reason:
3371                    "flexible(...) does not support sas/beta-logistic links; wiggle is only supported for jointly fit standard links"
3372                        .to_string(),
3373            });
3374        }
3375        return Ok(Some(LinkChoice {
3376            mode: LinkMode::Flexible,
3377            link,
3378            mixture_components: None,
3379        }));
3380    }
3381    if let Some(inner) = t
3382        .strip_prefix("blended(")
3383        .and_then(|s| s.strip_suffix(')'))
3384        .or_else(|| t.strip_prefix("mixture(").and_then(|s| s.strip_suffix(')')))
3385    {
3386        if flexible_flag {
3387            return Err(FormulaDslError::IncompatibleTerm {
3388                reason:
3389                    "--flexible-link cannot be combined with --link blended(...)/mixture(...); blended inverse links are not flexible-link mode"
3390                        .to_string(),
3391            });
3392        }
3393        let components = parse_link_component_list(inner)?;
3394        return Ok(Some(LinkChoice {
3395            mode: LinkMode::Strict,
3396            link: LinkFunction::Logit,
3397            mixture_components: Some(components),
3398        }));
3399    }
3400
3401    let link = parse_linkname(&t)?;
3402    if flexible_flag && !linkname_supports_joint_wiggle(link) {
3403        return Err(FormulaDslError::IncompatibleTerm {
3404            reason:
3405                "--flexible-link does not support sas/beta-logistic links; wiggle is only supported for jointly fit standard links"
3406                    .to_string(),
3407        });
3408    }
3409    Ok(Some(LinkChoice {
3410        mode: if flexible_flag {
3411            LinkMode::Flexible
3412        } else {
3413            LinkMode::Strict
3414        },
3415        link,
3416        mixture_components: None,
3417    }))
3418}
3419
3420pub fn parse_linkname(v: &str) -> Result<LinkFunction, FormulaDslError> {
3421    match v.trim() {
3422        "identity" => Ok(LinkFunction::Identity),
3423        "log" => Ok(LinkFunction::Log),
3424        "logit" | "binomial-logit" => Ok(LinkFunction::Logit),
3425        "probit" | "binomial-probit" => Ok(LinkFunction::Probit),
3426        "cloglog" | "binomial-cloglog" => Ok(LinkFunction::CLogLog),
3427        "loglog" => Ok(LinkFunction::LogLog),
3428        "cauchit" => Ok(LinkFunction::Cauchit),
3429        "sas" => Ok(LinkFunction::Sas),
3430        "beta-logistic" => Ok(LinkFunction::BetaLogistic),
3431        other => Err(FormulaDslError::UnknownIdentifier {
3432            reason: format!(
3433                "unsupported link type '{other}'; \
3434                 use one of identity|log|logit|probit|cloglog|loglog|cauchit|binomial-logit|binomial-probit|binomial-cloglog|sas|beta-logistic|blended(...)/mixture(...) or flexible(...). \
3435                 Both `--link <type>` (CLI flag) and `link(type=<type>)` (formula term) accept the same set."
3436            ),
3437        }),
3438    }
3439}
3440
3441pub fn parse_link_component(v: &str) -> Result<LinkComponent, String> {
3442    match v.trim() {
3443        "logit" => Ok(LinkComponent::Logit),
3444        "probit" => Ok(LinkComponent::Probit),
3445        "cloglog" => Ok(LinkComponent::CLogLog),
3446        "loglog" => Ok(LinkComponent::LogLog),
3447        "cauchit" => Ok(LinkComponent::Cauchit),
3448        other => Err(FormulaDslError::UnknownIdentifier {
3449            reason: format!(
3450                "unsupported blended-link component '{other}'; use probit|logit|cloglog|loglog|cauchit"
3451            ),
3452        }
3453        .into()),
3454    }
3455}
3456
3457pub fn parse_link_component_list(v: &str) -> Result<Vec<LinkComponent>, String> {
3458    let mut out = Vec::new();
3459    for part in v.split(',') {
3460        let trimmed = part.trim();
3461        if trimmed.is_empty() {
3462            continue;
3463        }
3464        let comp = parse_link_component(trimmed)?;
3465        if out.contains(&comp) {
3466            return Err(FormulaDslError::IncompatibleTerm {
3467                reason: "blended(...) cannot contain duplicate components".to_string(),
3468            }
3469            .into());
3470        }
3471        out.push(comp);
3472    }
3473    if out.len() < 2 {
3474        return Err(FormulaDslError::InvalidArgument {
3475            reason: "blended(...) requires at least two components".to_string(),
3476        }
3477        .into());
3478    }
3479    Ok(out)
3480}